Submission #1097385

#TimeUsernameProblemLanguageResultExecution timeMemory
1097385Alihan_8Mergers (JOI19_mergers)C++17
100 / 100
567 ms77400 KiB
#include <bits/stdc++.h>

using namespace std;

#define pb push_back
#define ar array

signed main(){
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr);
	
	int n, k; cin >> n >> k;
	
	vector <vector<int>> adj(n);
	
	for ( int i = 0; i + 1 < n; i++ ){
		int u, v; cin >> u >> v;
		
		--u, --v;
		
		adj[u].pb(v);
		adj[v].pb(u);
	}
	
	vector <int> s(n);
	
	for ( auto &x: s ){
		cin >> x; --x;
	}

	vector <int> tin(n), tout(n);
	vector <ar<int,2>> f(k, {n, -1});
	
	int ct = 0;
	
	auto trav = [&](auto trav, int u, int p) -> void{
		tin[u] = ++ct;
		
		if ( f[s[u]][0] == n ) f[s[u]][0] = ct;
		
		for ( auto &v: adj[u] ){
			if ( v != p ) trav(trav, v, u);
		} tout[u] = ct;
		
		f[s[u]][1] = ct;
	};
	
	trav(trav, 0, -1);
	
	vector <ar<int,2>> q(n);
	
	vector <int> bad(n);
	
	auto dfs = [&](auto dfs, int u, int p) -> void{
		q[u] = f[s[u]];
		
		for ( auto &v: adj[u] ){
			if ( v != p ){
				dfs(dfs, v, u);
				
				q[u][0] = min(q[u][0], q[v][0]);
				q[u][1] = max(q[u][1], q[v][1]);
			}
		}
		
		if ( tin[u] <= q[u][0] && tout[u] >= q[u][1] ){
			bad[u] = 1;
		}
	};
	
	dfs(dfs, 0, -1);
	
	vector <int> sub(n);
	
	bad[0] = 0;
	
	auto cals = [&](auto cals, int u, int p) -> void{
		sub[u] = bad[u];
		
		for ( auto &v: adj[u] ){
			if ( v != p ){
				cals(cals, v, u);
				
				sub[u] += sub[v];
			}
		}
	};
	
	cals(cals, 0, -1);
	
	int num = 0;
	
	for ( int u = 1; u < n; u++ ){
		if ( !bad[u] ) continue;
		
		if ( sub[u] == sub[0] || sub[u] == 1 ){
			num += 1;
		}
	}
	
	cout << (num + 1) / 2;
	
	cout << '\n';
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...