Submission #667322

#	Time	Username	Problem	Language	Result	Execution time	Memory
667322		dutinmeow	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	429 ms	39604 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

namespace std {

template<class Fun>
class y_combinator_result {
	Fun fun_;
public:
	template<class T>
	explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}

	template<class ...Args>
	decltype(auto) operator()(Args &&...args) {
		return fun_(std::ref(*this), std::forward<Args>(args)...);
	}
};

template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
	return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

} // namespace std

int main() {
	using namespace std;
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);

	int N, M;
	cin >> N >> M;
	vector<vector<int>> T(N);
	for (int i = 0; i < N - 1; i++) {
		int u, v;
		cin >> u >> v;
		u--, v--;
		T[u].push_back(v);
		T[v].push_back(u);
	}
	vector<int> A(N);
	for (int &a : A) {
		cin >> a;
		a--;
	} 

	vector<int> dep(N), hei(N);

	auto dfs1 = y_combinator([&](auto self, int u, int p) -> void {
		for (int v : T[u]) {
			if (v == p)
				continue;
			dep[v] = dep[u] + 1;
			self(v, u);
		}
	});

	dep[0] = 0;
	dfs1(0, -1);
	int ds = max_element(dep.begin(), dep.end()) - dep.begin();
	dfs1(ds, -1);
	int dt = max_element(dep.begin(), dep.end()) - dep.begin();
	swap(ds, dt);

	auto dfs2 = y_combinator([&](auto self, int u, int p) -> int {
		hei[u] = -1;
		for (int v : T[u]) {
			if (v == p)
				continue;
			dep[v] = dep[u] + 1;
			hei[u] = max(hei[u], self(v, u));
		}
		return ++hei[u];
	});

	int num = 0;
	vector<int> ans(N, 0), frq(M, 0);
	stack<int> stk;

	auto dfs3 = y_combinator([&](auto self, int u, int p) -> void {
		int h = 0, g = -1; 
		for (int v : T[u]) 
			if (v != p && hei[v] + 1 > h)
				h = hei[v] + 1, g = v;
		int s = 0; 
		for (int v : T[u])
			if (v != p && v != g)
				s = max(s, hei[v] + 1);	
		if (g != -1) {
			while (!stk.empty() && dep[stk.top()] >= dep[u] - s) {
				int v = stk.top(); stk.pop();
				if (--frq[A[v]] == 0)
					num--;
			}
			stk.push(u);
			if (frq[A[u]]++ == 0)
				num++;
			self(g, u);
		}

		while (!stk.empty() && dep[stk.top()] >= dep[u] - h) {
			int v = stk.top(); stk.pop();
			if (--frq[A[v]] == 0)
				num--;
		}

		ans[u] = max(ans[u], num);
		for (int v : T[u]) {
			if (v == p || v == g)
				continue;
			if (stk.empty() || stk.top() != u) {
				stk.push(u);
				if (frq[A[u]]++ == 0)
					num++;
			}
			self(v, u);	
		}
		if (!stk.empty() && stk.top() == u) {
			stk.pop();
			if (--frq[A[u]] == 0)
				num--;
		}
	});

	dep[ds] = 0;
	dfs2(ds, -1);
	dfs3(ds, -1);

	dep[dt] = 0;
	dfs2(dt, -1);
	dfs3(dt, -1);

	for (int x : ans)
		cout << x << '\n';
}

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