Submission #916647

#	Time	Username	Problem	Language	Result	Execution time	Memory
916647		a_l_i_r_e_z_a	Mergers (JOI19_mergers)	C++14	100 / 100	1010 ms	88972 KiB

mergers

// In the name oF the God 
#include <bits/stdc++.h>
#define ll long long
// #define int long long
#define pb push_back
#define F first
#define S second
#define mp make_pair
#define pii pair <int, int>
#define smin(x, y) (x) = min((x), (y))
#define smax(x, y) (x) = max((x), (y))
using namespace std;

const int inf = 1e9+7;
const int maxn = 5e5+5;

int n, k, how_many_col = 0, c[maxn], w[maxn], cnt[maxn], cnt_c[maxn];
bool kamel[maxn];
vector <pii> adj[maxn];
vector <int>  past;

void setW(int v, int p) {
    for (auto u : adj[v]) {
        if (u.F != p) {
            setW(u.F, v); w[v] += w[u.F];
        }
    }
    w[v]++;
}

void reset() {
    how_many_col = 0;
    for (int cc : past) cnt_c[cc] = 0;
    past.clear();
}

void DFS_(int v, int p) {
    for (auto u : adj[v]) {
        if (u.F != p ) {
            DFS_(u.F, v);
        }
    }
    if (cnt_c[c[v]] == 0) how_many_col++;
    cnt_c[c[v]]++;
    past.pb(c[v]);
    if (cnt[c[v]] == cnt_c[c[v]]) how_many_col--;
}

void DFS(int v, int p, int e) {
    int mx = 0, mxv = -1, mxe = -1;
    for (auto u : adj[v]) {
        if (u.F != p and w[u.F] > mx) {
            mx = w[u.F]; mxv = u.F; mxe = u.S;
        }
    }
    for (auto u : adj[v]) {
        if (u.F != p and u.F != mxv) {
            reset();
            DFS(u.F, v, u.S);
        }
    }
    if (mxv != -1) {
        reset();
        DFS(mxv, v, mxe);
        for (auto u : adj[v]) {
            if (u.F != p and u.F != mxv) {
                DFS_(u.F, v);
            }
        }
    }
    if (cnt_c[c[v]] == 0) how_many_col++;
    cnt_c[c[v]]++;
    past.pb(c[v]);
    if (cnt[c[v]] == cnt_c[c[v]]) how_many_col--;
    if (!how_many_col and e != -1) kamel[e] = 1;
}

bool seen[maxn];
pii ed[maxn];
int r = -1;
int amir = 0;

void DFS__(int v, int p){
    for(auto hh: adj[v]){
        int u = hh.F, e = hh.S;
        if(u == p) continue;
        DFS__(u, v);
        if(kamel[e]){
            seen[v] = 1;
            if(!seen[u]) amir++;
            seen[u] = 1;
        }
        seen[v] |= seen[u];
    }
    if(v == r){
        int x = 0;
        for(auto hh: adj[v]) x += seen[hh.F];
        if(x <= 1){
            amir++;
            // cout << "HI\n";
        }
    }
}

int32_t main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> k;
    for (int i = 0; i < n-1; i++) {
        int v, u; cin >> v >> u;
        adj[--v].pb(mp(--u, i));
        adj[u].pb(mp(v, i));
        ed[i] = mp(u, v);
    }
    for (int i = 0; i < n; i++) {
        cin >> c[i]; cnt[c[i]]++;
    }
    setW(0, -1);
    DFS(0, -1, -1);
    for(int i = 0; i < n - 1; i++){
        if(kamel[i]) r = ed[i].F;
    }
    if(r == -1){
        cout << 0 << '\n';
        exit(0);
    }
    DFS__(r, -1);
    cout << (amir + 1) / 2 << '\n';
    return 0;
}

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