Submission #1129151

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129151		khanhphucscratch	Capital City (JOI20_capital_city)	C++20	41 / 100	558 ms	122552 KiB

capital_city

#include<bits/stdc++.h>
using namespace std;
vector<int> ad[250005], color_list[250005], visited_node;
pair<int, int> sparse[500025][25];
int ans = 1e9, c[250005], h[250005], par[250005], sz[250005], tin[250005], dfsc = 0;
bool visited[250005], visited_color[250005], in_subtree[250005];
bool up_visited_in_subtree[250005], used_color_subtree[250005];
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;
}
void dfs_list_subtree(int u)
{
    //cout<<"B"<<u<<endl;
    visited_node.push_back(u);
    visited[u] = 1; sz[u] = 1;
    for(int v : ad[u]) if(visited[v] == 0 && visited_color[c[v]] == 0){
        dfs_list_subtree(v);
        sz[u] += sz[v];
    }
    visited[u] = 0;
}
int dfs_centroid_subtree(int u)
{
    for(int v : ad[u]) if(in_subtree[v] == 1 && sz[v] > sz[u]/2){
        sz[v] = sz[u] - sz[v]; swap(sz[u], sz[v]);
        return dfs_centroid_subtree(v);
    }
    return u;
}
void dfs_solve(int u)
{
    visited_node.clear();
    dfs_list_subtree(u);
    for(int i : visited_node) in_subtree[i] = 1;
    int r = dfs_centroid_subtree(u);
    //Now check if this color is good
    queue<int> w; w.push(c[r]); used_color_subtree[c[r]] = 1;
    int root_all = -1, curans = 0;
    bool ok = 1;
    //cout<<"A"<<u<<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]);
        }
        if(in_subtree[root_all] == 0){curans = 1e9; break;}
        //Up
        for(int i : use_node){
            int node = i;
            if(in_subtree[node] == 0){ok = 0; curans = 1e9; break;}
            if(node != root_all) node = par[node];
            while(node > 0){
                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];
            }
        }
        if(ok == 0) break;
    }
    vector<int> mylist = visited_node;
    ans = min(ans, curans); visited_color[c[r]] = 1;
    //cout<<"A"<<u<<" "<<r<<endl;
    for(int i : visited_node){
        in_subtree[i] = 0;
        up_visited_in_subtree[i] = 0;
        used_color_subtree[c[i]] = 0;
    }
    for(int i : visited_node) if(visited_color[c[i]] == 0) dfs_solve(i);
}
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);
    dfs_solve(1);
    cout<<ans-1;
}

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