Submission #1129139

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129139		cutimeo	Capital City (JOI20_capital_city)	C++20	1 / 100	3099 ms	97996 KiB

capital_city

//O(n^2)
#include<bits/stdc++.h>
using namespace std;
vector<int> ad[200005], color_list[200005], visited_node;
pair<int, int> sparse[400005][20];
int ans = 1e9, c[200005], h[200005], par[200005], tin[200005], dfsc = 0;
bool visited[200005], visited_color[200005];
bool up_visited_in_subtree[200005], used_color_subtree[200005];
void dfs_lca_original_tree(int u)
{
    visited[u] = 1; sparse[++dfsc][0] = {h[u], u}; tin[u] = dfsc;
    for(int v : ad[u]) if(visited[v] == 0){
        h[v] = h[u]+1; par[v] = u; dfs_lca_original_tree(v);
        sparse[++dfsc][0] = {h[u], u};
    }
    visited[u] = 0;
}
void preprocess_lca(int n)
{
    for(int j = 0; j < 19; j++){
        for(int i = 1; i <= n; i++){
            sparse[i][j+1] = sparse[i][j];
            if(i + (1 << j) <= n) sparse[i][j+1] = min(sparse[i][j], sparse[i+(1<<j)][j]);
        }
    }
}
int lca(int u, int v)
{
    int l = tin[u], r = tin[v];
    //cout<<"C"<<l<<" "<<r<<endl;
    if(l > r) swap(l, r);
    int x = __lg(r-l+1);
    return min(sparse[l][x], sparse[r-(1<<x)+1][x]).second;
}
int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    int n, k;
    cin>>n>>k;
    for(int i = 0; i < n-1; i++){
        int u, v; cin>>u>>v;
        ad[u].push_back(v);
        ad[v].push_back(u);
    }
    for(int i = 1; i <= n; i++){
        cin>>c[i];
        color_list[c[i]].push_back(i);
    }
    dfs_lca_original_tree(1); preprocess_lca(dfsc);
    for(int co = 1; co <= k; co++){
        for(int i = 1; i <= n; i++){
            up_visited_in_subtree[i] = 0;
            used_color_subtree[c[i]] = 0;
        }
        queue<int> w;
        w.push(co); used_color_subtree[co] = 1;
        int root_all = -1, curans = 0;
        //cout<<"A"<<co<<endl;
        while(w.size() > 0){
            int color = w.front(); w.pop(); curans++;
            vector<int> use_node = color_list[color];
            if(root_all != -1) use_node.push_back(root_all);
            //Calculate root_all
            root_all = use_node[0];
            //cout<<"B"<<color<<endl;
            for(int i = 1; i < use_node.size(); i++){
                root_all = lca(root_all, use_node[i]);
            }
            //Up
            for(int i : use_node){
                int node = i;
                while(node > 0){
                    //cout<<"C"<<node<<endl;
                    if(up_visited_in_subtree[node] == 1) break;
                    up_visited_in_subtree[node] = 1;
                    if(used_color_subtree[c[node]] == 0){
                        used_color_subtree[c[node]] = 1;
                        w.push(c[node]);
                    }
                    if(node == root_all) break;
                    node = par[node];
                }
            }
        }
        ans = min(ans, curans);
    }
    //dfs_solve(1);
    cout<<ans-1;
}

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