Submission #535352

#	Time	Username	Problem	Language	Result	Execution time	Memory
535352		79brue	Capital City (JOI20_capital_city)	C++14	100 / 100	1036 ms	194920 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

int n, k;
int arr[200002], idx[200002];
vector<int> link[200002];
vector<int> vec[200002];

int in[200002], inCnt, sz[200002];
int top[200002], depth[200002], par[200002], LCADB[200002][20];

void dfs_sz(int x, int p=-1){
    sz[x] = 1;
    if(p!=-1) link[x].erase(find(link[x].begin(), link[x].end(), p));
    for(auto &y: link[x]){
        depth[y] = depth[x] + 1;
        par[y] = LCADB[y][0] = x;
        dfs_sz(y, x);
        sz[x] += sz[y];
        if(sz[y] > sz[link[x][0]]) swap(y, link[x][0]);
    }
}

void dfs_hld(int x, int p=-1){
    in[x] = ++inCnt;
    idx[inCnt] = x;
    for(auto y: link[x]){
        top[y] = (y==link[x][0]) ? top[x] : y;
        dfs_hld(y, x);
    }
}

int getLCA(int x, int y){
    if(depth[x] < depth[y]) swap(x, y);
    for(int d=0; d<20; d++) if((depth[x] - depth[y]) & (1<<d)) x = LCADB[x][d];
    if(x==y) return x;
    for(int d=19; d>=0; d--) if(LCADB[x][d] != LCADB[y][d]) x = LCADB[x][d], y = LCADB[y][d];
    return par[x];
}

vector<int> graphLink[1000002];
vector<int> revLink[1000002];
bool visited[1000002];
stack<int> stk;
int ans = 1e9;
int cnt;
int sccNum[1000002];
vector<int> sccLink[1000002];
int weight[1000002];
int deg[1000002];

void sccDfs1(int x){
    visited[x] = 1;
    for(auto y: graphLink[x]){
        if(visited[y]) continue;
        sccDfs1(y);
    }
    stk.push(x);
}

void sccDfs2(int x, int t){
    visited[x] = 1;
    sccNum[x] = t;
    for(auto y: revLink[x]){
        if(visited[y]) continue;
        sccDfs2(y, t);
    }
}

void init(int i, int l, int r){
    if(l==r){
        graphLink[i].push_back(4*n+arr[idx[l]]);
        return;
    }
    int m = (l+r)>>1;
    graphLink[i].push_back(i*2);
    graphLink[i].push_back(i*2+1);
    init(i*2, l, m);
    init(i*2+1, m+1, r);
}

void update(int i, int l, int r, int s, int e, int x){
    if(r<s || e<l) return;
    if(s<=l && r<=e){
        graphLink[4*n+x].push_back(i);
        return;
    }
    int m = (l+r)>>1;
    update(i*2, l, m, s, e, x);
    update(i*2+1, m+1, r, s, e, x);
}

int main(){
    scanf("%d %d", &n, &k);
    for(int i=1; i<n; i++){
        int x, y;
        scanf("%d %d", &x, &y);
        link[x].push_back(y);
        link[y].push_back(x);
    }
    for(int i=1; i<=n; i++){
        scanf("%d", &arr[i]);
        vec[arr[i]].push_back(i);
    }

    dfs_sz(1);
    top[1] = 1; dfs_hld(1);
    for(int d=1; d<20; d++) for(int i=1; i<=n; i++) LCADB[i][d] = LCADB[LCADB[i][d-1]][d-1];
    init(1, 1, n);

    for(int i=1; i<=k; i++){
        sort(vec[i].begin(), vec[i].end(), [&](auto x, auto y){
            return in[x] < in[y];
        });
        for(int j=0; j<(int)vec[i].size() - 1; j++){
            int x = vec[i][j];
            int y = vec[i][j+1];
            int z = getLCA(x, y);

            while(top[x] != top[y]){
                if(depth[top[x]] < depth[top[y]]) swap(x, y);
                update(1, 1, n, in[top[x]], in[x], i);
                x = par[top[x]];
            }
            if(depth[x] > depth[y]) swap(x, y);
            update(1, 1, n, in[x], in[y], i);
        }
    }

//    for(int i=1; i<=5*n; i++){
//        printf("%d: ", i);
//        for(auto x: graphLink[i]) printf("%d ", x);
//        puts("");
//    }

    for(int i=1; i<=5*n; i++){
        if(!visited[i]) sccDfs1(i);
    }
    memset(visited, 0, sizeof(visited));
    for(int i=1; i<=5*n; i++) for(auto x: graphLink[i]) revLink[x].push_back(i);
    while(!stk.empty()){
        if(!visited[stk.top()]) sccDfs2(stk.top(), stk.top());
        stk.pop();
    }

    for(int i=1; i<=5*n; i++){
        for(auto x: graphLink[i]){
            if(sccNum[x] == sccNum[i]) continue;
            sccLink[sccNum[x]].push_back(sccNum[i]);
            deg[sccNum[i]]++;
        }
        if(4*n < i && i <= 4*n+k) weight[sccNum[i]]++;
    }

    queue<int> que;
    for(int i=1; i<=5*n; i++) if(i==sccNum[i] && !deg[i]) que.push(i);

    while(!que.empty()){
        int x = que.front(); que.pop();
        if(weight[x]) ans = min(weight[x]-1, ans);

        for(auto y: sccLink[x]){
            deg[y]--;
            weight[y] += weight[x];
            weight[y] = min(weight[y], 1000000);
            if(!deg[y]) que.push(y);
        }
    }

    printf("%d", ans);
}

Compilation message (stderr)

capital_city.cpp: In function 'int main()':
capital_city.cpp:121:17: warning: unused variable 'z' [-Wunused-variable]
  121 |             int z = getLCA(x, y);
      |                 ^
capital_city.cpp:97:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   97 |     scanf("%d %d", &n, &k);
      |     ~~~~~^~~~~~~~~~~~~~~~~
capital_city.cpp:100:14: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  100 |         scanf("%d %d", &x, &y);
      |         ~~~~~^~~~~~~~~~~~~~~~~
capital_city.cpp:105:14: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  105 |         scanf("%d", &arr[i]);
      |         ~~~~~^~~~~~~~~~~~~~~

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59