Submission #1370322

#TimeUsernameProblemLanguageResultExecution timeMemory
1370322sarahspeedyCapital City (JOI20_capital_city)C++20
0 / 100
3118 ms550200 KiB
#include<bits/stdc++.h>

#define int long long
#define pb push_back
#define vi vector<int>

using namespace std;

const int N = 2e5 + 5;

int n, k;

int c[N];

int sz[N];

int vis[N];

int activeNode[N];

int activeColor[N];

int parentArr[N];

int ans = 1e18;

vector<int> adj[N];

vector<int> nodesOfColor[N];

vector<int> componentNodes;

vector<int> usedColors;

void getSize(int node, int parent){

    sz[node] = 1;

    for(auto child : adj[node]){

        if(child == parent || vis[child]) continue;

        getSize(child, node);

        sz[node] += sz[child];
    }
}

int getCentroid(int node, int parent, int total){

    for(auto child : adj[node]){

        if(child == parent || vis[child]) continue;

        if(sz[child] > total / 2){

            return getCentroid(child, node, total);
        }
    }

    return node;
}

void dfsComponent(int node, int parent){

    parentArr[node] = parent;

    activeNode[node] = 1;

    componentNodes.pb(node);

    for(auto child : adj[node]){

        if(child == parent || vis[child]) continue;

        dfsComponent(child, node);
    }
}

void decompose(int entry){

    getSize(entry, -1);

    int centroid =
    getCentroid(entry, -1, sz[entry]);

    componentNodes.clear();

    dfsComponent(centroid, -1);

    queue<int> q;

    q.push(centroid);

    vector<int> colorUsed(k, 0);

    vector<int> nodeUsed;

    vector<int> colorList;

    int cntColors = 0;

    bool ok = true;

    activeColor[c[centroid]] = 1;

    colorList.pb(c[centroid]);

    cntColors++;

    while(!q.empty()){

        int node = q.front();
        q.pop();

        if(!activeNode[node]){

            ok = false;
            break;
        }

        if(!colorUsed[c[node]]){

            colorUsed[c[node]] = 1;

            for(auto u : nodesOfColor[c[node]]){

                if(!activeNode[u]){

                    ok = false;
                    break;
                }

                q.push(u);
            }

            if(!ok) break;

            cntColors++;
        }

        int cur = node;

        while(cur != -1){

            if(activeNode[cur]){

                q.push(cur);
            }

            cur = parentArr[cur];
        }
    }

    if(ok){

        ans = min(ans, cntColors - 1);
    }

    for(auto u : componentNodes){

        activeNode[u] = 0;
    }

    vis[centroid] = 1;

    for(auto child : adj[centroid]){

        if(!vis[child]){

            decompose(child);
        }
    }
}

signed main(){

    ios::sync_with_stdio(false);

    cin.tie(nullptr);

    cin >> n >> k;

    for(int i = 1; i < n; i++){

        int u, v;

        cin >> u >> v;

        u--;
        v--;

        adj[u].pb(v);

        adj[v].pb(u);
    }

    for(int i = 0; i < n; i++){

        cin >> c[i];

        c[i]--;

        nodesOfColor[c[i]].pb(i);
    }

    decompose(0);

    cout << ans << '\n';
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...