Submission #873074

#	Time	Username	Problem	Language	Result	Execution time	Memory
873074		serifefedartar	Capital City (JOI20_capital_city)	C++17	100 / 100	467 ms	41816 KiB

capital_city

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

#define fast ios::sync_with_stdio(0);cin.tie(0);
#define s second
#define f first
typedef long long ll;
const ll MOD = 1e9+9;
const ll LOGN = 20; 
const ll MAXN = 2e5 + 100;

vector<vector<int>> graph;
int city[MAXN], marked[MAXN], sz[MAXN], vis[MAXN], cVis[MAXN];
vector<int> seen, v[MAXN];
queue<int> q;

int get_sz(int node, int parent) {
	sz[node] = 1;
	for (auto u : graph[node]) {
		if (u == parent || marked[u])
			continue;
		sz[node] += get_sz(u, node);
	}
	return sz[node];
}

int find_centro(int node, int parent, int n) {
	for (auto u : graph[node]) {
		if (u != parent && !marked[u] && sz[u] * 2 >= n)
			return find_centro(u, node, n);
	}
	return node;
}

int par[MAXN], last[MAXN];
void dfs(int node, int parent, int c) {
	last[node] = c;
	vis[node] = false;
	seen.push_back(city[node]);
	for (auto u : graph[node]) {
		if (parent != u && !marked[u]) {
			par[u] = node;
			dfs(u, node, c);
		}
	}
}

int ans;
int cnt;
void decompose(int node) {
	int n = get_sz(node, node);
	int centro = find_centro(node, node, n);
	marked[centro] = true;
	last[centro] = centro; 

	q = queue<int>();
	seen.clear();
	seen.push_back(city[centro]); 
	for (auto u : graph[centro]) {
		if (!marked[u]) {
			par[u] = centro; 
			dfs(u, centro, centro);
		}
	}

	for (auto u : seen)
		cVis[u] = false;

	cnt = 0;
	q.push(city[centro]);
	bool control = true;
	while (!q.empty()) {
		int c = q.front();
		q.pop();

		if (cVis[c])
			continue;
		cVis[c] = true;
		cnt++;

		for (auto u : v[c]) {
			if (last[u] != centro) {
				control = false;
				break;
			}

			int now = u;
			while (now != centro) {
				if (vis[now])
					break;
				vis[now] = true; 
				q.push(city[now]);
				now = par[now];
			}
		}
	}
	if (control)
		ans = min(ans, cnt - 1);

	for (auto u : graph[centro]) {
		if (!marked[u])
			decompose(u);
	}
}

void dfs3(int node, int parent) {
	v[city[node]].push_back(node);
	for (auto u : graph[node]) {
		if (u == parent)
			continue;
		dfs3(u, node);
	}	
}

int main() {
	fast
	int N, K, A, B;
	cin >> N >> K;
	ans = N;

	graph = vector<vector<int>>(N+1, vector<int>());
	for (int i = 1; i < N; i++) {
		cin >> A >> B;
		graph[A].push_back(B);
		graph[B].push_back(A);
	}
	for (int i = 1; i <= N; i++)
		cin >> city[i];
	dfs3(1, 1);
	decompose(1);
	cout << ans << "\n";
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59