제출 #924090

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
924090		boris_mihov	수도 (JOI20_capital_city)	C++17	100 / 100	400 ms	45396 KiB

capital_city

#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>
#include <queue>

typedef long long llong;
const int MAXN = 200000 + 10;
const int INF  = 1e9;

int n, k;
int col[MAXN];
std::vector <int> g[MAXN];
std::vector <int> c[MAXN];
std::vector <int> inColor[MAXN];
int sz[MAXN];
int in[MAXN];
int out[MAXN];
int par[MAXN];
bool vis[MAXN];
int timer;

void calcSize(int node, int par)
{
    sz[node] = 1;
    for (const int &u : g[node])
    {
        if (u == par || vis[u])
        {
            continue;
        }

        calcSize(u, node);
        sz[node] += sz[u];
    }
}

int findCentroid(int node, int par, int globalSz)
{
    for (const int &u : g[node])
    {
        if (u != par && !vis[u] && sz[u] > globalSz / 2)
        {
            return findCentroid(u, node, globalSz);
        }
    }

    return node;
}

int decompose(int node)
{
    calcSize(node, 0);
    int cntr = findCentroid(node, 0, sz[node]);

    vis[cntr] = true;
    for (const int &u : g[cntr])
    {
        if (vis[u])
        {
            continue;
        }

        c[cntr].push_back(decompose(u));
    }
    
    return cntr;
}

int dep[MAXN];
void buildDFS(int node)
{
    in[node] = ++timer;
    for (const int &u : c[node])
    {
        dep[u] = dep[node] + 1;
        buildDFS(u);
    }

    out[node] = ++timer;
}

bool inSubtree(int x, int y)
{
    return in[y] <= in[x] && in[x] <= out[y];
}

int ans = INF;
void prepare(int node, int p, int lvl)
{
    par[node] = p;
    vis[col[node]] = false;
    for (const int &u : g[node])
    {
        if (u == p || dep[u] <= lvl)
        {
            continue;
        }

        prepare(u, node, lvl);
    }
}   

void solveDFS(int node)
{
    prepare(node, 0, dep[node]);
    std::queue <int> waitlist;
    waitlist.push(col[node]);
    vis[col[node]] = true;

    int res = -1;
    bool isBad = false;
    while (waitlist.size())
    {
        int top = waitlist.front();
        waitlist.pop();
        res++;

        for (const int &u : inColor[top])
        {
            if (!inSubtree(u, node))
            {
                isBad = true;
                break;
            }
            
            if (par[u] != 0 && !vis[col[par[u]]])
            {
                vis[col[par[u]]] = true;
                waitlist.push(col[par[u]]);
            }
        }

        if (isBad)
        {
            break;
        }
    }

    if (!isBad)
    {
        ans = std::min(ans, res);
    }

    for (const int &u : c[node])
    {
        solveDFS(u);
    }
}

void solve()
{
    int root = decompose(1);
    std::fill(vis + 1, vis + 1 + n, false);
    buildDFS(root);
    solveDFS(root);

    std::cout << ans << '\n';
}

void input()
{
    std::cin >> n >> k;
    for (int i = 1 ; i < n ; ++i)
    {
        int u, v;
        std::cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    for (int i = 1 ; i <= n ; ++i)
    {
        std::cin >> col[i];
        inColor[col[i]].push_back(i);
    }
}

void fastIOI()
{
    std::ios_base :: sync_with_stdio(0);
    std::cout.tie(nullptr);
    std::cin.tie(nullptr);
}

int main()
{
    fastIOI();
    input();
    solve();

    return 0;
}

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