Submission #1116854

#	Time	Username	Problem	Language	Result	Execution time	Memory
1116854		byebye75	Mergers (JOI19_mergers)	C++14	24 / 100	3078 ms	262144 KiB

mergers

#include <bits/stdc++.h>
#define int long long int
#define vii vector<int>
#define pii pair<int, int>
#define vpi vector<pii>
#define ff first
#define ss second

using namespace std;

const int N = 5e5 + 69;
const int K = 52;

int n, k;
int a[N], cnt[N], up[N], is_leaf[N];
vii g[N], c[K], t[N];
int have[N][K];

void dfs(int v, int p) {
    have[v][a[v]] = 1;
    for (auto u : g[v]) {
        if (u == p) continue;
        dfs(u, v);
        for (int j = 1; j <= k; j++) {
            have[v][j] += have[u][j];
        }
    }
    int ok = 1;
    for (int j = 1; j <= k; j++) {
        if (have[v][j] > 0)
            ok &= (have[v][j] == (int)c[j].size());
    }
    cnt[v] = ok;
}

void build(int v, int p, int c_node) {
    int x = c_node;
    if (cnt[v] && v > 1) {
        t[v].push_back(c_node);
        t[c_node].push_back(v);
        up[v] = c_node;
        x = v;
    }
    for (auto u : g[v]) {
        if (u == p) continue;
        build(u, v, x);
    }
}

void calc(int v, int p) {
    if (t[v].size() == 1 && v > 1)
        is_leaf[v] = 1;
    for (auto u : t[v]) {
        if (u == p) continue;
        calc(u, v);
        is_leaf[v] += is_leaf[u];
    }
}

signed main() {

    ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);

    cin >> n >> k;

    // Clear previous test case data
    for (int i = 1; i <= n; i++) {
        g[i].clear();
        t[i].clear();
        cnt[i] = is_leaf[i] = up[i] = 0;
        for (int j = 1; j <= k; j++) {
            have[i][j] = 0;
        }
    }
    for (int j = 1; j <= k; j++) {
        c[j].clear();
    }

    // Read tree edges
    for (int i = 1; i <= n - 1; i++) {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    // Read state assignments
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        c[a[i]].push_back(i);
    }

    dfs(1, 0);
    build(1, 0, 1);
    calc(1, 0);

    assert(cnt[1] == 1);

    int o = 0;
    for (int i = 1; i <= n; i++) {
        if (cnt[i] && t[i].size() == 1)
            o += 1;
    }

    int ans = (o + 1) / 2;
    cout << ans << '\n';

    return 0;
}

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