Submission #647412

#TimeUsernameProblemLanguageResultExecution timeMemory
647412ymmMergers (JOI19_mergers)C++17
100 / 100
1220 ms126520 KiB
#include <bits/stdc++.h>
#define Loop(x,l,r) for (ll x = (l); x < (ll)(r); ++x)
#define LoopR(x,l,r) for (ll x = (r)-1; x >= (ll)(l); --x)
typedef long long ll;
typedef std::pair<int, int> pii;
typedef std::pair<ll , ll > pll;
using namespace std;

const int N = 500'010;
int col[N];
int cnt_col[N];
bool is_in_sub[N];
vector<int> A[N];
int n, k;

class s2l {
private:
	int cnt_unfin = 0;
	map<int, int> mp;
public:
	void insert(int x, int cnt=1)
	{
		int &val = mp[x];
		cnt_unfin += val == 0;
		val += cnt;
		cnt_unfin -= val == cnt_col[x];
	}
	void merge(s2l *that)
	{
		if (this->mp.size() < that->mp.size()) {
			this->mp.swap(that->mp);
			swap(this->cnt_unfin, that->cnt_unfin);
		}
		for (auto [x, y] : that->mp)
			this->insert(x, y);
		delete(that);
	}
	bool has_unfin()
	{
		return cnt_unfin;
	}
};

s2l *dfs1(int v, int p)
{
	s2l *obj = new s2l;
	obj->insert(col[v]);
	for (int u : A[v]) {
		if (u != p)
			obj->merge(dfs1(u, v));
	}
	if (!obj->has_unfin() && p != -1) {
		is_in_sub[v] = 1;
		is_in_sub[p] = 1;
	}
	return obj;
}

bool dfs2(int v, int p)
{
	for (int u : A[v]) {
		if (u != p)
			if (dfs2(u, v))
				is_in_sub[v] = 1;
	}
	return is_in_sub[v];
}

int dfs3(int v, int p)
{
	int nbr = 0;
	int leaf = 0;
	for (int u : A[v]) {
		if (!is_in_sub[u])
			continue;
		nbr++;
		if (u == p)
			continue;
		leaf += dfs3(u, v);
	}
	leaf += nbr == 1;
	return leaf;
}

int main()
{
	cin.tie(0) -> sync_with_stdio(false);
	cin >> n >> k;
	Loop (i,1,n) {
		int v, u;
		cin >> v >> u;
		--v; --u;
		A[v].push_back(u);
		A[u].push_back(v);
	}
	Loop (i,0,n) {
		int c;
		cin >> c;
		col[i] = c-1;
		++cnt_col[c-1];
	}
	dfs1(0, -1); // I CAN LEAK ANYTHING!
	int rt;
	for (rt = 0; rt < n && !is_in_sub[rt]; ++rt);
	if (rt == n) {
		cout << "0\n";
		return 0;
	}
	dfs2(rt, -1);
	int leaves = dfs3(rt, -1);
	cout << (leaves+1)/2 << '\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...