Submission #1116862

#	Time	Username	Problem	Language	Result	Execution time	Memory
1116862		byebye75	Mergers (JOI19_mergers)	C++14	48 / 100	177 ms	262144 KiB

mergers

#include <bits/stdc++.h>
//#include "bigint.h"
#define db(x) cerr << #x << ": " << x << endl
#define print cerr << "Ah shit, here we go again" << endl
#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 = 1e5 + 10;
const int INF = 1e9 + 10;

int n, k;
int a[N], c[N], sub[N], ans, b[N], root, cnt[N];
vii g[N];
vector<int> mark[N];

void dfs(int v, int p) {
    sub[v] = 1;
    mark[v][a[v]] = 1;
    for (auto u : g[v]) {
        if (u == p) continue;
        dfs(u, v);
        sub[v] += sub[u];
        cnt[v] += cnt[u];
        for (int i = 1; i <= k; i++) {
            mark[v][i] |= mark[u][i];
        }
    }
    int x = 0;
    for (int y = 1; y <= k; y++) {
        if (mark[v][y]) x += c[y];
    }
    if ((x == sub[v] && v != root)) {
        b[v] = 1;
        ++cnt[v];
    }
}

void dfs1(int v, int p) {
    if (b[v] == 1) {
        if (cnt[v] == 1 || cnt[v] == cnt[root]) ++ans;
    }
    for (auto u : g[v]) {
        if (u == p) continue;
        dfs1(u, v);
    }
}



signed main() {

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

    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int v, u; cin >> v >> u;
        g[v].push_back(u);
        g[u].push_back(v);
    }
    for (int i = 1; i <= n; i++) {
        mark[i].resize(k + 1, 0);
    }
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        ++c[a[i]];
    }
    root = 1;
    dfs(root, 0);
    dfs1(root, 0);
    cout << (ans + 1) / 2 << '\n';

    return 0;
}

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