Submission #217116

#	Time	Username	Problem	Language	Result	Execution time	Memory
217116		anonymous	Capital City (JOI20_capital_city)	C++14	100 / 100	766 ms	32372 KiB

capital_city

#include<iostream>
#include<set>
#include<vector>
#include<cassert>
#define MAXN 200005
using namespace std;
int N, K, best = 1<<30, C[MAXN];
vector <int> adj[MAXN], type[MAXN]; //adj and towns in each city
int cmp[MAXN], id = 1;
//Solve for a node
int par[MAXN];
bool done[MAXN];
vector <int> S;
set <int> Used;
void reroot(int u, int prev) {
    if (cmp[u] != cmp[prev]) {return;}
    par[u] = prev, done[u]=false;
    for (auto v: adj[u]) {if (v != prev) {reroot(v, u);}}
}

int slv(int v) { //centre v
    int ans = 0;
    S.clear();
    Used.clear();
    reroot(v, v);
    Used.insert(C[v]);
    par[v] = 0;
    for (auto k: type[C[v]]) {
        if (cmp[k] != cmp[v]) {
            return(1<<30);
        } else {
            S.push_back(k);
        }
    }
    while (S.size()) {
        int cur = S.back();
        S.pop_back();
        while (cur != 0 && !done[cur]) {
          done[cur] = true;
          if (Used.find(C[cur]) == Used.end()) {
            for (auto k: type[C[cur]]) {
                if (cmp[k] != cmp[cur]) {
                    return(1<<30);
                } else {
                    S.push_back(k);
                }
            }
            Used.insert(C[cur]);
          }
          cur = par[cur];
        }
    }
    return(Used.size());
}

//Centroid decomp
int Size[MAXN];
int resize(int u, int prev) {
     if (cmp[u] != cmp[prev]) {return(0);}
     Size[u] = 1;
     for (auto v: adj[u]) {
        if (v != prev) {Size[u] += resize(v,u);}
     }
     //printf("Size[%d] = %d\n",u,Size[u]);
     return(Size[u]);
}

int FindCentroid(int rt, int u, int prev) {
     if (cmp[u] != cmp[prev]) {return(-1);}
     int MaxSz = 0, SumSz=1;
     for (auto v: adj[u]) {
        if (v != prev && cmp[v] == cmp[u]) {
            MaxSz = max(MaxSz, Size[v]);
            SumSz += Size[v];
        }
     }
     //printf("%d %d %d\n", Size[rt], MaxSz, SumSz);
     if (2*MaxSz <= Size[rt]+1 && 2*SumSz + 1 >= Size[rt]) {
        return(u);
     }
     for (auto v: adj[u]) {
        if (v == prev) {continue;}
        int res = FindCentroid(rt, v, u);
        if (res > 0) {return(res);}
     }
     return(-1);
}

void relabel(int u, int prev, int cmp_prev) {
    if (cmp[u] != cmp_prev) {return;}
    cmp[u] = id;
    for (auto v: adj[u]) {
        if (v != prev) {relabel(v,u,cmp_prev);}
    }
}

void decomp(int u) {
    //printf("Decomp Node %d: cmp %d\n", u,cmp[u]);
    resize(u, u);
    int cent = Size[u] > 1 ? FindCentroid(u, u, u) : u;
    //printf("Chosen Centroid: %d\n",cent);
    assert(cent > 0);
    best = min(best, slv(cent));
    //cmp[cent] = -1;
    for (auto v: adj[cent]) {
        if (cmp[v] != cmp[cent]) {continue;}
        relabel(v, cent, cmp[cent]);
        id++;
        decomp(v);
    }
}

int main() {
    //freopen("capin.txt","r",stdin);
    scanf("%d %d", &N, &K);
    for (int i=1; i<N; i++) {
        int a,b;
        scanf("%d %d",&a,&b);
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    for (int i=1; i<=N; i++) {
        scanf("%d",&C[i]);
        type[C[i]].push_back(i);
    }
    decomp(1);
    printf("%d\n",best-1);
}

Compilation message (stderr)

capital_city.cpp: In function 'int slv(int)':
capital_city.cpp:22:9: warning: unused variable 'ans' [-Wunused-variable]
     int ans = 0;
         ^~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:115:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     scanf("%d %d", &N, &K);
     ~~~~~^~~~~~~~~~~~~~~~~
capital_city.cpp:118:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%d %d",&a,&b);
         ~~~~~^~~~~~~~~~~~~~~
capital_city.cpp:123:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%d",&C[i]);
         ~~~~~^~~~~~~~~~~~

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