Submission #462589

#	Time	Username	Problem	Language	Result	Execution time	Memory
462589		arwaeystoamneg	Capital City (JOI20_capital_city)	C++17	100 / 100	1079 ms	44244 KiB

capital_city

// EXPLOSION!
#define _CRT_SECURE_NO_WARNINGS
#include<bits/stdc++.h>
#include<unordered_set>
#include<unordered_map>
#include<chrono>

using namespace std;
typedef pair<int, int> pii;
typedef long long ll;
typedef pair<ll, ll> pll;
typedef long double ld;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<pair<int, int>> vpi;
typedef vector<pair<ll, ll>> vpll;

#define FOR(i,a,b) for (int i = (a); i < (b); ++i)
#define F0R(i,a) FOR(i,0,a)
#define ROF(i,a,b) for (int i = (b)-1; i >= (a); --i)
#define R0F(i,a) ROF(i,0,a)
#define trav(a,x) for (auto& a: x)

#define pb push_back
#define mp make_pair
#define rsz resize
#define sz(x) int(x.size())
#define all(x) x.begin(),x.end()
#define f first
#define s second
#define cont continue
#define endl '\n'
//#define ednl '\n'
#define test int testc;cin>>testc;while(testc--)
#define pr(a, b) trav(x,a)cerr << x << b; cerr << endl;
#define message cout << "Hello World" << endl;
const int dx[4] = { 1,0,-1,0 }, dy[4] = { 0,1,0,-1 }; // for every grid problem!!
const ll linf = 4000000000000000000LL;
const ll inf = 1000000007;//998244353    

void pv(vi a) { trav(x, a)cout << x << " "; cout << endl; }void pv(vll a) { trav(x, a)cout << x << " "; cout << endl; }void pv(vector<vi>a) {
	F0R(i, sz(a)) { cout << i << endl; pv(a[i]); cout << endl; }
}void pv(vector<vll>a) { F0R(i, sz(a)) { cout << i << endl; pv(a[i]); }cout << endl; }void pv(vector<string>a) { trav(x, a)cout << x << endl; cout << endl; }
void setIO(string s) {
	ios_base::sync_with_stdio(0); cin.tie(0);
#ifdef arwaeystoamneg
	if (sz(s))
	{
		freopen((s + ".in").c_str(), "r", stdin);
		if (s != "test1")
			freopen((s + ".out").c_str(), "w", stdout);
	}
#endif
}
int n, k, c, ans;
vector<vi>adj, cadj;
vi visited, sizes, parent, co, depth, up, v;
struct centroid
{
	void dfs(int i, int p)
	{
		// visit i
		sizes[i] = 1;
		for (auto x : adj[i])
		{
			if (x == p || visited[i] == 1)
				continue;
			dfs(x, i);
			sizes[i] += sizes[x];
		}
	}
	int dfs_centroid(int i, int p, int n)
	{
		// visit i, find the centroid c
		// recurrent if i has a subtree larger than n/2;
		for (auto x : adj[i])
		{
			if (x == p)
				continue;
			if (sizes[x] > n / 2)
				return dfs_centroid(x, i, n);
		}
		return i;
	}
	void build(int i, int p)
	{
		// first dfs to generate size of subtree starting from node i
		// 2nd dfs to find the centroid
		dfs(i, p);
		int centroid = dfs_centroid(i, p, sizes[i]);
		visited[centroid] = 1;
		parent[centroid] = p;
		//cout << centroid << endl;
		// remove the edge from centroid to all its child, build from its child
		for (auto x : adj[centroid])
		{
			if (x == parent[centroid])
				continue;
			//cout << "working on subtree of " << x << endl;
			if (visited[x])
				continue;
			build(x, centroid);
		}
	}
	void init()
	{
		n = sz(adj);
		visited.rsz(n);
		sizes.rsz(n);
		parent.rsz(n);
		build(0, -1);
	}
};

vi a;
void read()
{
	cin >> n >> k;
	adj.rsz(n);
	F0R(i, n - 1)
	{
		int x, y;
		cin >> x >> y;
		x--, y--;
		adj[x].pb(y);
		adj[y].pb(x);
	}
	a.rsz(n);
	trav(x, a)cin >> x, x--;
}
void prep()
{
	co.rsz(k);
	F0R(i, n)co[a[i]]++;
}
void cent()
{
	centroid g;
	g.init();
	cadj.rsz(n);
	F0R(i, n)if (parent[i] != -1)cadj[parent[i]].pb(i); else c = i;
}
vector<vi>b;
vi toclear;
void init(int i, int p, int d)
{
	toclear.pb(a[i]);
	v[i] = 0;
	up[i] = p;
	b[a[i]].pb(i);
	trav(x, adj[i])
	{
		if (x == p || depth[x] < d)continue;
		init(x, i, d);
	}
}
void dfs_centroid(int i, int p, int d)
{
	depth[i] = d;
	trav(x, cadj[i])
	{
		if (x == p)continue;
		dfs_centroid(x, i, d + 1);
	}
	toclear.clear();
	init(i, -1, d);
	int cur = 0;
	vi todo = { a[i] };
	v[i] = 1;
	while (sz(todo))
	{
		int x = todo.back();
		todo.pop_back();
		if (sz(b[x]) == 0)continue;
		cur++;
		if (co[x] > sz(b[x])) { cur = inf; break; }
		trav(y, b[x])
		{
			int t = y;
			while (!v[t])
			{
				todo.pb(a[t]);
				v[t] = 1;
				t = up[t];
			}
		}
		b[x].clear();
	}
	trav(x, toclear)b[x].clear();
	ans = min(ans, cur - 1);
}
int main()
{
	setIO("test1");
	read();
	prep();
	cent();
	depth.rsz(n);
	up.rsz(n);
	v.rsz(n);
	b.rsz(k);
	ans = inf;
	dfs_centroid(c, -1, 0);
	cout << ans << endl;
}

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