Submission #330891

#	Time	Username	Problem	Language	Result	Execution time	Memory
330891		limabeans	Mergers (JOI19_mergers)	C++17	100 / 100	1498 ms	137580 KiB

mergers

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

template<typename T>
void out(T x) { cout << x << endl; exit(0); }
#define watch(x) cout << (#x) << " is " << (x) << endl





using ll = long long;


const int maxn = 5e5 + 10;



int n, k;
vector<int> g[maxn];
int a[maxn];


const int LOG = 20;
int dep[maxn];
int par[LOG+1][maxn];

void dfs(int at, int p) {
    for (int j=1; j<LOG; j++) {
	par[j][at] = par[j-1][par[j-1][at]];
    }
    
    for (int to: g[at]) {
	if (to == p) continue;
	par[0][to] = at;
	dep[to] = 1+dep[at];
	dfs(to, at);
    }
}

int lca(int u, int v) {
    if (dep[u]>dep[v]) {
	swap(u,v);
    }
    // u
    // v

    int dx = dep[v]-dep[u];
    for (int j=LOG-1; j>=0; j--) {
	if (dx>>j&1) {
	    v = par[j][v];
	}
    }

    if (u==v) return u;

    for (int j=LOG-1; j>=0; j--) {
	if (par[j][u] != par[j][v]) {
	    u = par[j][u];
	    v = par[j][v];
	}
    }

    return par[0][v];
}

vector<int> nodes[maxn];
int dp[maxn];


void dfs2(int at, int p) {
    for (int to: g[at]) {
	if (to == p) continue;
	dfs2(to, at);
	dp[at] += dp[to];
    }
}


int deg[maxn];
int P[maxn];


void dfs3(int at, int p) {
    if (dp[at]==0) {
	if (p!=-1) {
	    // add edge
	    // p-->at
	    if (at != p) {
		deg[at]++;
		deg[P[p]]++;
	    }
	    P[at] = at;
	}
    } else {
	P[at] = P[p];
    }
    
    for (int to: g[at]) {
	if (to == p) continue;
	dfs3(to, at);
    }
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0);  cout.tie(0);

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

    for (int i=0; i<n; i++) {
	cin>>a[i];
	--a[i];
	nodes[a[i]].push_back(i);
    }

    dfs(0, -1);

    for (int i=0; i<n; i++) {
	dp[i] = 1;
    }

    for (int j=0; j<k; j++) {
	int mid = nodes[j][0];
	for (int x: nodes[j]) {
	    mid = lca(mid, x);
	}
	dp[mid] -= nodes[j].size();
    }


    dfs2(0, -1);

    // for (int i=0; i<n; i++) {
    // 	cout<<i<<": "<<dp[i]<<endl;
    // }


    dfs3(0, -1);

    int leaves = 0;
    for (int i=0; i<n; i++) {
	if (deg[i]==1) leaves++;
    }



    cout<<(leaves+1)/2<<endl;
    return 0;
}

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