Submission #1275299

#TimeUsernameProblemLanguageResultExecution timeMemory
1275299kaiboyUnique Cities (JOI19_ho_t5)C++20
100 / 100
248 ms58068 KiB
#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

const int N = 200000;

vector<int> ej[N], ej_[N];
int aa[N], dds[N], ddt[N], *dd, d_, i_;
int hh[N], ii_[N], cc_[N];
int ii[N], pp[N], qq[N], rr[N], n_;
int kk[N], cc[N];

void dfs_dd(int p, int i, int d) {
	if (d_ < d)
		d_ = d, i_ = i;
	dd[i] = d++;
	for (int j : ej[i])
		if (j != p)
			dfs_dd(i, j, d);
}

int dfs_hh(int p, int i) {
	int h = 0;
	for (int j : ej[i])
		if (j != p)
			h = max(h, dfs_hh(i, j) + 1);
	return hh[i] = h;
}

int search(int i_, int d) {
	while (i_ && dds[ii[i_]] >= d)
		i_ = qq[i_] && dds[ii[qq[i_]]] >= d ? qq[i_] : pp[i_];
	return i_;
}

void dfs(int p, int i, int i_) {
	if (dds[i] >= ddt[i])
		ii_[i] = search(i_, dds[i] - hh[i]);
	int h0 = 0, h1 = 0;
	for (int j : ej[i])
		if (j != p) {
			int h = hh[j] + 1;
			if (h0 < h)
				h1 = h0, h0 = h;
			else if (h1 < h)
				h1 = h;
		}
	for (int j : ej[i])
		if (j != p) {
			int p_ = search(i_, dds[i] - (hh[j] + 1 == h0 ? h1 : h0));
			ii[n_] = i, ej_[pp[n_] = p_].push_back(n_);
			if (rr[p_] == rr[qq[p_]])
				qq[n_] = qq[qq[p_]], rr[n_] = rr[p_] + 1;
			else
				qq[n_] = p_, rr[n_] = 0;
			dfs(i, j, n_++);
		}
}

void dfs_kk(int i) {
	static int c;
	cc[i] = c += i && !kk[aa[ii[i]]]++;
	for (int j : ej_[i])
		dfs_kk(j);
	c -= i && !--kk[aa[ii[i]]];
}

int main() {
	ios_base::sync_with_stdio(false), cin.tie(NULL);
	int n, a_; cin >> n >> a_;
	for (int h = 0; h < n - 1; h++) {
		int i, j; cin >> i >> j, i--, j--;
		ej[i].push_back(j), ej[j].push_back(i);
	}
	for (int i = 0; i < n; i++)
		cin >> aa[i], aa[i]--;
	dd = dds, d_ = -1, dfs_dd(-1, 0, 0); int s = i_;
	dd = dds, d_ = -1, dfs_dd(-1, s, 0); int t = i_;
	dd = ddt, dfs_dd(-1, t, 0);
	for (int z = 0; z < 2; z++) {
		dfs_hh(-1, s);
		n_ = 1, dfs(-1, s, 0);
		dfs_kk(0);
		for (int i = 0; i < n; i++)
			if (dds[i] >= ddt[i])
				cc_[i] = cc[ii_[i]];
		swap(s, t);
		for (int i = 0; i < n; i++)
			swap(dds[i], ddt[i]), ej_[i].clear();
	}
	for (int i = 0; i < n; i++)
		cout << cc_[i] << '\n';
	return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...