Submission #1249055

#	Time	Username	Problem	Language	Result	Execution time	Memory
1249055		vlomaczk	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	438 ms	32844 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
typedef long long ll;
using namespace __gnu_pbds;
using namespace std;

template <typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

int x, y;
int M = 200'010;
vector<vector<int>> g(M);
vector<int> c(M), d(M), h(M), colors(M), ans(M);
int cnt=0;
stack<int> stos;

void add(int v) {
	stos.push(v);
	colors[c[v]]++;
	if(colors[c[v]]==1) cnt++;
}

void pop() {
	if(stos.empty()) return;
	int v = stos.top();
	stos.pop();
	colors[c[v]]--;
	if(colors[c[v]]==0) cnt--;
}

void dfs1(int v, int p) {
	d[v] = d[p]+1;
	h[v] = 0;
	for(auto u : g[v]) {
		if(u==p) continue;
		dfs1(u, v);
		h[v] = max(h[v], h[u]+1);
	}
}

void dfs2(int v, int p) {
	pair<int, int> maks1 = {0, 0};
	pair<int, int> maks2 = {0, 0};
	for(auto u : g[v]) if(u!=p) maks1 = max(maks1, {h[u]+1, u});
	for(auto u : g[v]) if(u!=p && u!=maks1.second) maks2 = max(maks2, {h[u]+1, u});
	if(maks1.second) {
		while(stos.size() && d[stos.top()] >= d[v]-maks2.first) pop();
		add(v);
		dfs2(maks1.second, v);
		if(stos.size() && stos.top()==v) pop();
		while(stos.size() && d[stos.top()] >= d[v]-maks1.first) pop();
		for(auto u : g[v]) {
			if(u==p || u==maks1.second) continue;
			add(v);
			dfs2(u, v);
			if(stos.size() && stos.top()==v) pop();
		} 
	}
	ans[v] = max(ans[v], cnt);
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);

	int n, k;
	cin >> n >> k;
	for(int i=1; i<n; ++i) {
		int a, b;
		cin >> a >> b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	for(int i=1; i<=n; ++i) cin >> c[i];
	dfs1(1, 0);
	pair<int, int> maks = {0, 0};
	for(int i=1; i<=n; ++i) maks = max(maks, {d[i], i});
	int L = maks.second;
	dfs1(L, 0);
	dfs2(L, 0);
	maks = {0, 0};
	for(int i=1; i<=n; ++i) maks = max(maks, {d[i], i});
	int R = maks.second;
	dfs1(R, 0);
	dfs2(R, 0);
	for(int i=1; i<=n; ++i) cout << ans[i] << " "; cout << "\n";

	return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32