Submission #729981

#	Time	Username	Problem	Language	Result	Execution time	Memory
729981		pavement	Unique Cities (JOI19_ho_t5)	C++17	32 / 100	1227 ms	250136 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int long long
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define ppb pop_back
#define eb emplace_back
#define g0(a) get<0>(a)
#define g1(a) get<1>(a)
#define g2(a) get<2>(a)
#define g3(a) get<3>(a)
#define g4(a) get<4>(a)
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
using db = double;
using ll = long long;
using ld = long double;
using ii = pair<int, int>;
using iii = tuple<int, int, int>;
using iiii = tuple<int, int, int, int>;
using iiiii = tuple<int, int, int, int, int>;
template<class key, class value = null_type, class cmp = less<key> >
using ordered_set = tree<key, value, cmp, rb_tree_tag, tree_order_statistics_node_update>;

int N, M, lon[200005], ans_st[200005], ans_ed[200005], lon_st[200005], lon_ed[200005], dist_st[200005], dist_ed[200005], C[200005];
ii m = mp(-1, -1);
vector<int> adj[200005];
stack<int> rep[200005];

struct node {
	node *left, *right;
	int S, E, pv;
	ii val;
	void combine() {
		if (S == E) return;
		if (left->val.first == right->val.first) val = mp(left->val.first, left->val.second + right->val.second);
		else val = min(left->val, right->val);
	}
	node(int _s, int _e) : S(_s), E(_e), pv(0) {
		if (S == E) {
			val = mp(0, 1);
			return;
		}
		int M = (S + E) >> 1;
		left = new node(S, M);
		right = new node(M + 1, E);
		combine();
	}
	void prop() {
		if (S == E || !pv) return;
		left->val.first += pv;
		left->pv += pv;
		right->val.first += pv;
		right->pv += pv;
		pv = 0;
	}
	void upd(int l, int r, int v) {
		if (l > E || r < S) return;
		if (l <= S && E <= r) {
			val.first += v;
			pv += v;
			return;
		}
		prop();
		left->upd(l, r, v);
		right->upd(l, r, v);
		combine();
	}
	ii qry(int l, int r) {
		if (l > E || r < S) return mp((int)1e16, 0);
		if (l <= S && E <= r) return val;
		prop();
		auto lq = left->qry(l, r), rq = right->qry(l, r);
		if (lq.first == rq.first) return mp(lq.first, lq.second + rq.second);
		else return min(lq, rq);
	}
} *root;

void dfs(int n, int e = -1, int d = 0) {
	m = max(m, mp(d, n));
	for (auto u : adj[n]) if (u != e) {
		dfs(u, n, d + 1);
	}
}

void init(int n, int lon[], int dist[], int e = -1) {
	for (auto u : adj[n]) if (u != e) {
		dist[u] = dist[n] + 1;
		init(u, lon, dist, n);
		lon[n] = max(lon[n], lon[u]);
	}
	lon[n]++;
}

void get_ans(int n, int ans[], int lon[], int e = -1, int d = 0) {
	bool pushed = 0;
	if (rep[C[n]].empty() || root->qry(rep[C[n]].top(), rep[C[n]].top()).first) {
		rep[C[n]].push(d);
		pushed = 1;
	} else {
		root->upd(d, d, 1);
	}
	root->upd(d - lon[n] + 1, d - 1, 1);
	auto tmp = root->qry(0, d);
	ans[n] = (tmp.first ? 0ll : tmp.second - pushed);
	root->upd(d - lon[n] + 1, d - 1, -1);
	multiset<int> m;
	for (auto u : adj[n]) if (u != e) {
		m.insert(lon[u]);
	}
	for (auto u : adj[n]) if (u != e) {
		m.erase(m.find(lon[u]));
		int x = (m.empty() ? 0ll : *m.rbegin());
		root->upd(d - x, d - 1, 1);
		get_ans(u, ans, lon, n, d + 1);
		root->upd(d - x, d - 1, -1);
		m.insert(lon[u]);
	}
	if (pushed) {
		assert(!rep[C[n]].empty() && rep[C[n]].top() == d);
		rep[C[n]].pop();
	} else {
		root->upd(d, d, -1);
	}
}

main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> N >> M;
	for (int i = 1, u, v; i < N; i++) {
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for (int i = 1; i <= N; i++) {
		cin >> C[i];
	}
	dfs(1);
	int st = m.second;
	m = mp(-1, -1);
	dfs(st);
	int ed = m.second;
	root = new node(0, N - 1);
	init(st, lon_st, dist_st);
	get_ans(st, ans_st, lon_st);
	root = new node(0, N - 1);
	init(ed, lon_ed, dist_ed);
	for (int i = 1; i <= M; i++) {
		while (!rep[i].empty()) rep[i].pop();
	}
	get_ans(ed, ans_ed, lon_ed);
	for (int i = 1; i <= N; i++) {
		if (dist_st[i] <= dist_ed[i]) {
			cout << ans_ed[i] << '\n';
		} else {
			cout << ans_st[i] << '\n';
		}
	}
}

Compilation message (stderr)

joi2019_ho_t5.cpp:130:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
  130 | main() {
      | ^~~~

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