Submission #900614

#	Time	Username	Problem	Language	Result	Execution time	Memory
900614		LucaIlie	Capital City (JOI20_capital_city)	C++17	100 / 100	589 ms	43680 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

const int MAX_N = 2e5 + 1;
const int MAX_K = 2e5 + 1;
bool isCentroid[MAX_N + 1], cityIn[MAX_N + 1];
int city[MAX_N + 1], sizee[MAX_N + 1], frecv[MAX_N + 1], frecvC[MAX_K + 1], parent[MAX_N + 1];
vector<int> adj[MAX_N + 1], towns[MAX_K + 1];

void calcSizes( int u, int p ) {
    sizee[u] = 1;
    for ( int v: adj[u] ) {
        if ( v == p || isCentroid[v] )
            continue;
        calcSizes( v, u );
        sizee[u] += sizee[v];
    }
}

int centroid, sizetot;
void findCentroid( int u, int p ) {
    int maxSize = sizetot - sizee[u];
    for ( int v: adj[u] ) {
        if ( v == p || isCentroid[v] )
            continue;
        findCentroid( v, u );
        maxSize = max( maxSize, sizee[v] );
    }
    if ( maxSize <= sizetot / 2 )
        centroid = u;
}

vector<int> vert;
void dfs( int u, int p ) {
    parent[u] = p;
    vert.push_back( u );
    frecvC[city[u]]++;
    for ( int v: adj[u] ) {
        if ( v == p || isCentroid[v] )
            continue;
        dfs( v, u );
    }
}

int minMerges;
void decomp( int r ) {
    calcSizes( r, 0 );
    sizetot = sizee[r];
    findCentroid( r, 0 );
    int c = centroid;

    vert.clear();
    dfs( c, 0 );

    queue<int> q;
    bool sePoate = true;
    int m = 0;
    if ( frecvC[city[c]] == frecv[city[c]] ) {
        cityIn[city[c]] = true;
        for ( int v: towns[city[c]] )
            q.push( v );
        while ( !q.empty() ) {
            int v = q.front();
            q.pop();

            int x = city[parent[v]];
            if ( x != 0 && !cityIn[x] ) {
                cityIn[x] = true;
                m++;
                if ( frecvC[x] != frecv[x] ) {
                    sePoate = false;
                    break;
                }
                for ( int v: towns[x] )
                    q.push( v );
            }
        }
    } else
        sePoate = false;

    if ( sePoate )
        minMerges = min( minMerges, m );

    for ( int v: vert ) {
        frecvC[city[v]] = 0;
        cityIn[city[v]] = false;
    }

    isCentroid[c] = true;
    for ( int v: adj[c] ) {
        if ( !isCentroid[v] )
            decomp( v );
    }
}

int main() {
    int n, k;

    cin >> n >> k;
    for ( int i = 0; i < n - 1; i++ ) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back( v );
        adj[v].push_back( u );
    }
    for ( int v = 1; v <= n; v++ ) {
        cin >> city[v];
        frecv[city[v]]++;
        towns[city[v]].push_back( v );
    }

    minMerges = k - 1;
    decomp( 1 );

    cout << minMerges;

    return 0;
}

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