Submission #416566

#	Time	Username	Problem	Language	Result	Execution time	Memory
416566		amirmohammad_nezami	Capital City (JOI20_capital_city)	C++14	100 / 100	1399 ms	42396 KiB

capital_city

//                                                          In the name of God 
#include <bits/stdc++.h>
using namespace std;
#define F first
#define S second
#define PB push_back
#define MP make_pair
#define lc 2 * v
#define rc 2 * v + 1
#define mid (s + e) / 2
#define _sz(e) (int)e.size()
#define _all(e) e.begin() , e.end()
#define _lcm(a , b) (1ll * a * b) / __gcd(a , b)
#define ll long long
#define int long long
#define ld long double
#define pii pair <int , int>
#define pcc pair <char , char>
#define pll pair <long long , long long>
#define piii pair <int , pair <int , int>>
#define FAST ios::sync_with_stdio(0);cin.tie(0)
#pragma GCC optimize("Ofast")
// #pragma GCC target("ve,ve2,ve3,vse3,ve4,avmini,avmini2")

const int maxn = 3e5 + 10 , N = 2e5 + 10 , SQ = 320 , base = 800287 , mod = 1e9 + 7 , INF = 1e14 + 10 , lg = 22;
const ld eps = 1e-4;

int n , k , c[maxn] , res[maxn] , sz[maxn] , par[maxn]; 
bool death[maxn] , sub[maxn] , mark[maxn] , visited[maxn];
vector <int> edges[maxn] , town[maxn] , tmp;

void dfs_sz(int v , int p = -1) {
	sz[v] = 1;
	for (auto u : edges[v]) {
		if(u != p && !death[u]) {
			dfs_sz(u , v);
			sz[v] += sz[u];
		}
	}
}

void dfs_sub(int v , int p = -1) {
	sub[v] = 1 , par[v] = p , tmp.PB(v);
	for (auto u : edges[v]) {
		if(u != p && !death[u]) {
			dfs_sub(u , v);
		}
	}
}

int find_centroid(int v , int tot , int p = -1) {
	for (auto u : edges[v]) {
		if(u != p && !death[u] && sz[u] > tot / 2) {
			return find_centroid(u , tot , v);
		}
	}
	return v;
}

void calc(int x) {
	dfs_sub(x);
	queue <int> q; q.push(x);

	while(_sz(q)) {
		int v = q.front(); q.pop();
		if(visited[v]) {continue;}
		
		visited[v] = 1 , tmp.PB(v);

		if(!mark[c[v]]) {
			for (auto u : town[c[v]]) {
				if(!sub[u] || death[u]) {
					res[x] = INF;
					return ;
				}
				if(!visited[u]) {
					q.push(u);
				}
			}
			res[x]++;
			mark[c[v]] = 1 , tmp.PB(c[v]);
		}
		
		if(par[v] != -1 && !visited[par[v]]) {
			q.push(par[v]);
		}
	}
}

void centroid_tree() {
	queue <int> q; dfs_sz(0);
	q.push(find_centroid(0 , sz[0]));

	while(_sz(q)) {
		int v = q.front(); q.pop();

		calc(v);
		for (auto u : tmp) {
			sub[u] = mark[u] = visited[u] = 0;
		}
		tmp.clear(); 
		death[v] = 1;
	
		for (auto u : edges[v]) {
			if(!death[u]) {
				dfs_sz(u);
				q.push(find_centroid(u , sz[u]));
			}
 		}
	}
}

int32_t main() {	
	FAST;
	cin >> n >> k;
	for (int i = 0; i < n - 1; ++i) {
		int x , y; cin >> x >> y;
		x-- , y--;
		edges[x].PB(y);
		edges[y].PB(x);
	}
	for (int i = 0; i < n; ++i) {
		cin >> c[i]; c[i]--;
		town[c[i]].PB(i);
	}

	centroid_tree();

	int ans = k - 1;
	for (int i = 0; i < n; ++i) {
		ans = min(ans , res[i] - 1);
	}

	cout << ans << '\n';
} 
//  Mistakes:
//  * Read the problem carefully.
//  * Check maxn.
//  * Overflows.

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