Submission #1129009

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129009		OI_Account	Mergers (JOI19_mergers)	C++20	100 / 100	638 ms	104692 KiB

mergers

#include <bits/stdc++.h>
using namespace std;

const int N = 500'000;

int n, k, s[N + 10], par[N + 10];
int tim, st[N + 10], ft[N + 10];
int mnPack[N + 10], mxPack[N + 10];
int mn[N + 10], mx[N + 10], deg[N + 10];
vector<int> adj[N + 10], pack[N + 10];
vector<pair<int, int>> edge;

void readInput() {
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 1; i <= n; i++) {
        cin >> s[i];
        pack[s[i]].push_back(i);
    }
    fill(par + 1, par + n + 1, -1);
}

void dfs(int u = 1, int par = 0) {
    st[u] = ++tim;
    for (auto v: adj[u])
        if (v != par)
            dfs(v, u);
    ft[u] = tim;
}

void calcMinMax() {
    for (int i = 1; i <= k; i++) {
        mnPack[i] = N + 10;
        mxPack[i] = -1;
        for (auto u: pack[i]) {
            mnPack[i] = min(mnPack[i], st[u]);
            mxPack[i] = max(mxPack[i], st[u]);
        }
    }
}

int getPar(int u) {
    return (par[u] < 0? u: (par[u] = getPar(par[u])));
}

bool merg(int u, int v) {
    if ((u = getPar(u)) == (v = getPar(v)))
        return false;
    if (par[v] < par[u])
        swap(u, v);
    par[u] += par[v];
    par[v] = u;
    return true;
}

void dfsEdge(int u = 1, int par = 0) {
    mn[u] = mnPack[s[u]];
    mx[u] = mxPack[s[u]];
    for (auto v: adj[u])
        if (v != par) {
            dfsEdge(v, u);
            mn[u] = min(mn[u], mn[v]);
            mx[u] = max(mx[u], mx[v]);
        }
    if (par) {
        if (st[u] <= mn[u] && mx[u] <= ft[u])
            edge.push_back({u, par});
        else
            merg(u, par);
    }
    //cout << u << ' ' << par << ": " << st[u] << ' ' << ft[u] << " - " << mn[u] << mx[u] << endl;
}

void calcDeg() {
    for (auto [u, v]: edge) {
        deg[getPar(u)]++;
        deg[getPar(v)]++;
    }
}

void calcAns() {
    int cntLeaf = 0;
    for (int i = 1; i <= n; i++)
        if (i == getPar(i) && deg[i] == 1)
            cntLeaf++;
    cout << (cntLeaf + 1) / 2 << flush;
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    readInput();
    dfs();
    calcMinMax();
    dfsEdge();
    calcDeg();
    calcAns();
    return 0;
}

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30