Submission #1085889

#	Time	Username	Problem	Language	Result	Execution time	Memory
1085889		underwaterkillerwhale	Mergers (JOI19_mergers)	C++17	100 / 100	652 ms	74832 KiB

mergers

#include <bits/stdc++.h>
#define ll              long long
#define pii             pair<int,int>
#define pll             pair<ll,ll>
#define rep(i,m,n)      for(int i=(m); i<=(n); i++)
#define reb(i,m,n)      for(int i=(m); i>=(n); i--)
#define iter(id, v)     for(auto id : v)
#define fs              first
#define se              second
#define MP              make_pair
#define pb              push_back
#define bit(msk, i)     ((msk >> i) & 1)
#define SZ(v)           (ll)v.size()
#define ALL(v)          v.begin(),v.end()

using namespace std;

mt19937_64 rd(chrono :: steady_clock :: now ().time_since_epoch().count());
ll Rand (ll l, ll r) { return uniform_int_distribution<ll> (l, r) (rd); }

const int N = 5e5 + 7;
const int Mod = 1e9 + 2022;
const int INF = 1e9;
const ll BASE = 137;
const int szBL = 350;

int n, K;
vector<int> ke[N];
int a[N];
int ccnt[N], num_cpn, tot_cpn[N], num_ok;
int ein[N], nC[N];
pii dp[N];
int res = 0;

void pdfs (int u, int p) {
    static int time_dfs = 0;
    ein[u] = ++time_dfs;
    nC[u] = 1;
    iter (&v, ke[u]) {
        if (v != p) {
            pdfs(v, u);
            nC[u] += nC[v];
        }
    }
}

void Assign (int u, int delta) {
    if (delta == 1) {
        if (ccnt[a[u]] == 0) ++num_cpn;
        ccnt[a[u]]++;
        if (ccnt[a[u]] == tot_cpn[a[u]]) ++num_ok;
    }
    else {
        if (ccnt[a[u]] == tot_cpn[a[u]]) --num_ok;
        ccnt[a[u]]--;
        if (ccnt[a[u]] == 0) --num_cpn;
    }
}

void update (int u, int p, int delta) {
    Assign(u, delta);
    iter (&v, ke[u]) {
        if (v != p) update (v, u, delta);
    }
}

void dfs (int u, int p) {
    int mxV = -1;
    iter (&v, ke[u]) if (v != p && (mxV == -1 || nC[v] > nC[mxV])) mxV = v;
    iter (&v, ke[u]) if (v != p && v != mxV) {
        dfs(v, u);
        update (v, u, -1);
    }
    if (mxV != -1) dfs(mxV, u);
    iter (&v, ke[u]) if (v != p && v != mxV) {
        update(v, u, 1);
    }
    int rem = 0;
    int numP = 0;
    iter (&v, ke[u]) {
        if (v != p) {
            numP += min(rem, dp[v].fs);
            rem = abs(rem - dp[v].fs);
        }
    }
    Assign (u, 1);
    bool okE = 0;
    if (u != 1 && num_ok == num_cpn) okE = 1;
    if (numP == 0) {
        if (rem == 0) dp[u] = {okE, 0};
        else {
            iter (&v, ke[u]) if (v != p && dp[v].fs != 0) {
                dp[u] = dp[v];
            }
            dp[u].se = dp[u].se | (okE & (rem >= 2));
        }
        return;
    }
    if (rem == 0 && numP > 0) {
        --numP;
        rem = 2;
    }
    res += numP;
    dp[u] = {rem, okE & (rem >= 2)};
    assert(rem <= 2);
}

void solution() {
    cin >> n >> K;
    rep (i, 1, n - 1) {
        int u, v;
        cin >> u >> v;
        ke[u].pb(v);
        ke[v].pb(u);
    }
    rep (i, 1, n) cin >> a[i], tot_cpn[a[i]]++;
    pdfs(1, 0);
    dfs(1, 0);
    if (dp[1].fs <= 1) res += dp[1].fs;
    else res += 1 + dp[1].se;
    cout << res <<"\n";
}

#define file(name) freopen(name".inp","r",stdin); \
freopen(name".out","w",stdout);
int main () {
//    file("c");
    ios_base :: sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int num_Test = 1;
//    cin >> num_Test;
    while (num_Test--)
        solution();
}
/*
no bug +5
9 6
7 6
9 1
2 4
4 5
9 2
8 6
9 3
8 1
3
1
6
4
4
2
5
5
6


*/

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