Submission #473640

#	Time	Username	Problem	Language	Result	Execution time	Memory
473640		iulia13	Capital City (JOI20_capital_city)	C++14	100 / 100	1025 ms	43492 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;
const int N = 2e5 + 5;
vector <int> g[N];
int color[N], viz[N], sz[N], team[N], newDad[N], usedColor[N], newViz[N];
vector <int> col[N];
int n, k, ans = N;
void dfsSize(int nod, int dad, int boss)
{
    sz[nod] = 1;
    newViz[nod] = 0;
    team[nod] = boss;
    newDad[nod] = dad;
    usedColor[color[nod]] = 0;
    for (auto son : g[nod])
    {
        if (son == dad || viz[son])
            continue;
        dfsSize(son, nod, boss);
        sz[nod] += sz[son];
    }
}
int findCentroid(int nod, int dad, int myN)
{
    for (auto son : g[nod])
    {
        if (son == dad || viz[son])
            continue;
        if (sz[son] > myN / 2)
            return findCentroid(son, nod, myN);
    }
    return nod;
}
void centroidDecomp(int nod)
{
    dfsSize(nod, 0, nod);
    int cnt = 0, mid = findCentroid(nod, 0, sz[nod]), ok = 1;
    dfsSize(mid, 0, mid);
    viz[mid] = 1;
    queue <int> q;
    q.push(mid);
    while (ok && !q.empty())
    {
        int x = q.front();
        int cx = color[x];
        q.pop();
        if (usedColor[cx])
            continue;
        cnt++;
        usedColor[cx] = 1;
        for (auto y : col[cx])
        {
            if (team[y] != mid)
            {
                ok = 0;
                break;
            }
            int cy = color[newDad[y]];
            if (y == mid || usedColor[cy])
                continue;
            q.push(newDad[y]);
        }
    }
    if (ok)
    ans = min(ans, cnt - 1);
    for (auto z : g[mid])
        if (!viz[z])
            centroidDecomp(z);
}
int main()
{
    //freopen("x.in", "r", stdin);
  //  freopen("x.out", "w", stdout);
    cin >> n >> k;
    for (int i = 1; i < n; i++)
    {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    for (int i = 1; i <= n; i++)
    {
        cin >> color[i];
        col[color[i]].push_back(i);
    }
    centroidDecomp(1);
    cout << ans;
    return 0;
}

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