Submission #1249036

#	Time	Username	Problem	Language	Result	Execution time	Memory
1249036		vlomaczk	Unique Cities (JOI19_ho_t5)	C++20	0 / 100	226 ms	17112 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 M = 200'010;
vector<vector<int>> g(M);
vector<int> c(M);
int x, y;
vector<int> distX(M), distY(M), hX(M), hY(M);
vector<bool> closer_to_X(M);
vector<int> ans(M), dist(M);
int odp = 0;
vector<int> path;
int max_dist = 0;
int how_many_colors = 0;
vector<int> color_count(M);
 
void dfs(int v, int p) {
	dist[v] = dist[p]+1;
	if(dist[v] > dist[max_dist]) max_dist=v;
	for(auto u : g[v]) if(u!=p) dfs(u, v);
}
 
pair<int, int> get_max_dist(int v) {
	dist.assign(M, -1);
	max_dist=0;
	dfs(v, 0);
	return {max_dist, dist[max_dist]};
}

void dfsX(int v, int p) {
	distX[v] = distX[p] + 1;
	for(auto u : g[v]) {
		if(u==p) continue;
		dfsX(u, v);
		hX[v] = max(hX[v], hX[u]+1);
	}
}

void dfsY(int v, int p) {
	distY[v] = distY[p] + 1;
	for(auto u : g[v]) {
		if(u==p) continue;
		dfsY(u, v);
	}
}

void insert(int v) {
	color_count[c[v]]++;
	if(color_count[c[v]]==1) how_many_colors++;
}

void erase(int v) {
	color_count[c[v]]--;
	if(color_count[c[v]]==0) how_many_colors--;
}

void path_dfsX(int v, int p) {
	if(closer_to_X[v]==0) ans[v] = how_many_colors;
	map<int, int> h;
	for(auto u : g[v]) if(u!=p) h[hX[u]]++;
	path.push_back(v);
	for(auto u : g[v]) {
		if(u==p) continue;
		if(g[u].size()==1) insert(v);
		if(h[hX[u]]==1) {
			int idx = path.size()-hX[u]-2;
			if(idx >= 0) insert(path[idx]);
		}
		path_dfsX(u, v);
		if(g[u].size()==1) erase(v);
		if(h[hX[u]]==1) {
			int idx = path.size()-hX[u]-2;
			if(idx >= 0) erase(path[idx]);
		}
	}
	path.pop_back();
}

void path_dfsY(int v, int p) {
	if(closer_to_X[v]) ans[v] = how_many_colors;
	map<int, int> h;
	for(auto u : g[v]) if(u!=p) h[hY[u]]++;
	path.push_back(v);
	for(auto u : g[v]) {
		if(u==p) continue;
		if(g[u].size()==1) insert(v);
		if(h[hY[u]]==1) {
			int idx = path.size()-hY[u]-2;
			if(idx >= 0) insert(path[idx]);
		}
		path_dfsY(u, v);
		if(g[u].size()==1) erase(v);
		if(h[hY[u]]==1) {
			int idx = path.size()-hY[u]-2;
			if(idx >= 0) erase(path[idx]);
		}
	}
	path.pop_back();
}

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];
	int x = get_max_dist(1).first;
	int y = get_max_dist(x).first;
	distX[0] = distY[0] = -1;
	dfsX(x, 0); dfsY(y, 0);
	for(int i=1; i<=n; ++i) closer_to_X[i] = (distX[i] < distY[i]);
	path_dfsX(x, 0); path_dfsY(y, 0);
	for(int i=1; i<=n; ++i) cout << ans[i] << " "; cout << "\n";

	return 0;
}

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