Submission #934354

#	Time	Username	Problem	Language	Result	Execution time	Memory
934354		penguin133	Mergers (JOI19_mergers)	C++17	100 / 100	1104 ms	223436 KiB

mergers

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

#define int long long
#define pi pair<int, int>
#define pii pair<int, pi>
#define fi first
#define se second
#ifdef _WIN32
#define getchar_unlocked _getchar_nolock
#endif
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

int n, k;
int S[500005], A[500005];
vector <int> adj[500005];
int deg[500005];
vector <int> shit[500005];
map <int, int> mp[500005];

pi dp[500005];

void dfs(int x, int par){
	mp[x][S[x]] = 1;
	int cnt = 0;
	for(auto i : adj[x]){
		if(i == par)continue;
		dfs(i, x);
		cnt += (mp[i].empty());
		if(shit[x].size() < shit[i].size())swap(shit[x], shit[i]);
		for(auto j : shit[i])shit[x].push_back(j);
		dp[x].fi += dp[i].fi;
		dp[x].se += dp[i].se;
		if(dp[x].se >= 3){
			dp[x].fi += dp[x].se / 2, dp[x].se %= 2;
			if(!dp[x].se)dp[x].se = 2, dp[x].fi--;
		}
		if(mp[x].size() < mp[i].size())swap(mp[x], mp[i]);
		for(auto [j, k] : mp[i]){
			mp[x][j] += k;
			if(mp[x][j] == A[j])mp[x].erase(j);
		}
	}
	if(mp[x].find(S[x]) != mp[x].end() && mp[x][S[x]] == A[S[x]])mp[x].erase(S[x]);
	if(mp[x].empty()){
		deg[x] += shit[x].size();
		for(auto i : shit[x])deg[i]++;
		shit[x] = {x};
	}
	/*
	cout << x << '\n';
	for(auto [i, j] : mp[x])cout << i << ' ' << j << '\n';
	cout << '\n';
	*/
}

void solve(){
	cin >> n >> k;
	for(int i = 1; i < n; i++){
		int a, b; cin >> a >> b;
		adj[a].push_back(b); adj[b].push_back(a);
	}
	for(int i = 1; i <= n; i++)cin >> S[i], A[S[i]]++;
	dfs(1, -1);
	//cout << dp[1].fi + dp[1].se;
	//cout << dp[1].fi + (dp[1].se + 1) / 2;
	int ans = 0;
	for(int i = 1; i <= n; i++)ans += (deg[i] == 1);
	cout << (ans + 1) / 2;
}

main(){
	ios::sync_with_stdio(0);cin.tie(0);
	int tc = 1;
	//cin >> tc;
	for(int tc1=1;tc1<=tc;tc1++){
		// cout << "Case #" << tc1 << ": ";
		solve();
	}
}

Compilation message (stderr)

mergers.cpp:72:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   72 | main(){
      | ^~~~

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