Submission #1352603

#TimeUsernameProblemLanguageResultExecution timeMemory
1352603Jawad_Akbar_JJCapital City (JOI20_capital_city)C++20
100 / 100
515 ms121632 KiB
#include <iostream>
#include <vector>
#include <set>
#include <algorithm>

using namespace std;
const int N = 1<<18, M = N * 3;
vector<int> nei[M], adj[M], jda[M], lst[M], scc, ord;
int pr[M], ch[M], col[M], id[M], num[M], tp[M], lca[M], vis[M], lab[M], inv[M], cur;

void Add(int u, int v){
	if (v == 0)
		v = N + N + 1;
	adj[u].push_back(v);
	jda[v].push_back(u);
}

void build(int cr = 1){
	lab[cr] = cr + N + N;
	if (cr >= N){
		lab[cr] = inv[cr - N];
		return;
	}

	int lc = cr + cr, rc = lc + 1;
	build(rc);
	build(lc);

	Add(lab[cr], lab[lc]);
	Add(lab[cr], lab[rc]);
}

void rangeAdd(int l, int r, int x){
	for (l += N - 1, r += N + 1;l + 1 < r; l /= 2, r /= 2){
		if (!(l & 1))
			Add(x, lab[l + 1]);
		if (r & 1)
			Add(x, lab[r - 1]);
	}
}

void dfs1(int u){
	ch[u] = 1;
	for (int &i : nei[u]){
		nei[i].erase(find(begin(nei[i]), end(nei[i]), u));
		dfs1(i);
		pr[i] = u;
		ch[u] += ch[i];
		if (ch[i] > ch[nei[u][0]])
			swap(nei[u][0], i);
	}
}

void dfs2(int u, int t){
	num[u] = ++cur;
	inv[cur] = u;
	tp[u] = t;

	for (int i : nei[u])
		dfs2(i, (i == nei[u][0] ? t : i));
}

int LCA(int u, int v){
	for (; tp[u] != tp[v]; v = pr[tp[v]]){
		if (num[v] < num[u])
			swap(u, v);
	}
	if (num[v] < num[u])
		swap(u, v);
	return u;
}

void dfs3(int u){
	int v = u, t = lca[col[u]];
	while (tp[v] != tp[t]){
		rangeAdd(num[tp[v]], num[v], u);
		v = pr[tp[v]];
	}
	rangeAdd(num[t], num[v], u);

	for (int i : nei[u])
		dfs3(i);
}

void dfs4(int u){
	vis[u] = 1;
	for (int i : adj[u])
		if (vis[i] != 1)
			dfs4(i);
	ord.push_back(u);
}

void dfs5(int u){
	vis[u] = 2;
	scc.push_back(u);
	id[u] = cur;
	for (int i : jda[u])
		if (vis[i] != 2)
			dfs5(i);
}


int main(){
	ios_base::sync_with_stdio(false); cin.tie(NULL);
	int n, k;
	cin>>n>>k;

	for (int i=1;i<n;i++){
		int a, b;
		cin>>a>>b;
		nei[a].push_back(b);
		nei[b].push_back(a);
	}

	dfs1(1);
	dfs2(1, 1);

	for (int i=1;i<=n;i++){
		cin>>col[i];
		lst[col[i]].push_back(i);
		if (lca[col[i]] == 0)
			lca[col[i]] = i;
		else
			lca[col[i]] = LCA(i, lca[col[i]]);
	}
	for (int i=1;i<=k;i++){
		int L = lst[i].back();
		for (int j : lst[i])
			Add(L, j), L = j;
	}

	build();

	dfs3(1);

	for (int i=1;i<=n;i++)
		if (vis[i] == 0)
			dfs4(i);
	if (!vis[N + N + 1])
		dfs4(N + N + 1);

	int ans = n;
	while (ord.size() > 0){
		int u = ord.back(); ord.pop_back();
		if (vis[u] != 2){
			++cur;
			scc = {};
			dfs5(u);

			set<int> st = {0};
			bool t = 1;
			for (int i : scc){
				for (int j : adj[i]){
					if (id[j] != cur)
						t = 0;
				}
				st.insert(col[i]);
			}
			if (t)
				ans = min(ans, (int)st.size());
		}
	}
	cout<<max(0, ans - 2)<<endl;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...