Submission #252029

#	Time	Username	Problem	Language	Result	Execution time	Memory
252029		Bruteforceman	Capital City (JOI20_capital_city)	C++11	100 / 100	1720 ms	56820 KiB

capital_city

#include "bits/stdc++.h"
using namespace std;
using pii = pair <int, int>;
vector <int> v[200005];
vector <int> g[200005];
int c[200005];

int disc[200005], fin[200005], dep[200005];
int anc[20][200010];
const int logn = 19;
int cur;
bool del[200005];
int n, k;

int lca(int x, int y) {
	if(dep[x] > dep[y]) {
		swap(x, y);
	}
	for(int i = logn; i >= 0; i--) {
		if(dep[y] - (1 << i) >= dep[x]) {
			y = anc[i][y];
		}
	}
	if(x == y) return x;
	for(int i = logn; i >= 0; i--) {
		if(anc[i][x] != anc[i][y]) {
			x = anc[i][x];
			y = anc[i][y];
		}
	}
	return anc[0][x];
}
void dfs(int x, int par) {
	anc[0][x] = par;
	disc[x] = ++cur;
	for(int i = 1; i <= logn; i++) {
		anc[i][x] = anc[i - 1][anc[i - 1][x]];
	}
	for(auto i : g[x]) {
		if(del[i]) continue;
		if(i - par) {
			dep[i] = 1 + dep[x];
			dfs(i, x);
		}
	}
	fin[x] = cur;
}

int par[200005];
int root(int x) {
	if(par[x] == x) return x;
	return par[x] = root(par[x]);
}
void join(int x, int y) {
	x = root(x);
	y = root(y);
	if(x != y) {
		par[x] = y;
	}
}

bool done[200005];
int ans;

int solve(int col) {
	queue <int> Q;
	set <int> cont;
	cont.insert(col);
	Q.push(col);	
	while(!Q.empty()) {
		int x = Q.front();
		Q.pop();
		vector <int> u = v[x];
		sort(u.begin(), u.end(), [&] (int p, int q) { return disc[p] < disc[q]; });
		for(int i = 1; i < v[x].size(); i++) {
			u.push_back(lca(u[i - 1], u[i]));
		}
		sort(u.begin(), u.end(), [&] (int p, int q) { return disc[p] < disc[q]; });
		stack <int> st;
		int last = -1;
		for(int i : u) {
			if(i == last) continue;
			while(!st.empty() && fin[st.top()] < disc[i]) {
				st.pop();
			}
			if(!st.empty()) {
				int cur = i;
				while(dep[st.top()] < dep[cur]) {
					if(done[c[cur]]) return k;
					if(!cont.count(c[cur])) Q.push(c[cur]);
					cont.insert(c[cur]);
					join(cur, anc[0][cur]);
					cur = root(anc[0][cur]);
				}
			}
			else {
				if(done[c[i]]) return k;
				if(!cont.count(c[i])) Q.push(c[i]);
				cont.insert(c[i]);
				st.push(i);
			}
			last = i; 
		}
	}
	return cont.size() - 1;
}

int sub[200010], opt[200010];
void subtree(int x, int par) {
	sub[x] = 1;
	for(auto i : g[x]) {
		if(del[i]) continue;
		if(i - par) {
			subtree(i, x);
			sub[x] += sub[i];
		}
	}
}
int centroid(int x, int par, int range) {
	for(auto i : g[x]) {
		if(del[i]) continue;
		if(i - par) {
			if(sub[i] > range) return centroid(i, x, range);
		}
	}
	return x;
}
void process(int x, int p, int root) {
	if(opt[c[x]] == -1) {
		opt[c[x]] = root;
	} else if (opt[c[x]] != root) {
		done[c[x]] = true;
	}
	for(auto i : g[x]) {
		if(del[i]) continue;
		if(i - p) {
			process(i, x, root);
		}
	}
}
void reset(int x, int p) {
	par[x] = x;
	opt[c[x]] = -1;
	for(auto i : g[x]) {
		if(del[i]) continue;
		if(i - p) {
			reset(i, x);
		}
	}
} 
void decomp(int x) {
	subtree(x, 0);
	x = centroid(x, 0, sub[x] / 2);	
	reset(x, 0);
	del[x] = true;
	if(!done[c[x]]) {
		ans = min(ans, solve(c[x]));
		done[c[x]] = true;
	}
	for(int i : g[x]) {
		if(!del[i]) {
			process(i, x, i);
		}
	}
	for(auto i : g[x]) {
		if(!del[i]) {
			decomp(i);
		}
	}
}

int main(int argc, char const *argv[])
{
	scanf("%d %d", &n, &k);
	for(int i = 1; i < n; i++) {
		int p, q;
		scanf("%d %d", &p, &q);
		g[p].emplace_back(q);
		g[q].emplace_back(p);
	}
	for(int i = 1; i <= n; i++) {
		scanf("%d", &c[i]);
		v[c[i]].push_back(i);
	}
	dfs(1, 0);
	ans = k;
	decomp(1);
	printf("%d\n", ans);
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int solve(int)':
capital_city.cpp:75:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i = 1; i < v[x].size(); i++) {
                  ~~^~~~~~~~~~~~~
capital_city.cpp: In function 'int main(int, const char**)':
capital_city.cpp:174:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d %d", &n, &k);
  ~~~~~^~~~~~~~~~~~~~~~~
capital_city.cpp:177:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d", &p, &q);
   ~~~~~^~~~~~~~~~~~~~~~~
capital_city.cpp:182:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d", &c[i]);
   ~~~~~^~~~~~~~~~~~~

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