Submission #405463

#	Time	Username	Problem	Language	Result	Execution time	Memory
405463		flappybird	Unique Cities (JOI19_ho_t5)	C++14	4 / 100	2092 ms	141140 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#include <unordered_map>
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx2")
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pll;
#define MAX 201010
#define MAXS 20
#define INF 1000000000000000001
#define bb ' '
#define ln '\n'
struct segtree {
	ll N;
	ll s;
	vector<ll> tree, l, r;
	void update(ll x, ll a) {
		x += s - 1;
		tree[x] = a;
		x /= 2;
		while (x) tree[x] = max(tree[x * 2], tree[x * 2 + 1]), x /= 2;
	}
	ll query(ll low, ll high, ll loc = 1) {
		if (l[loc] == low && r[loc] == high) return tree[loc];
		if (r[loc * 2] >= high) return query(low, high, loc * 2);
		if (l[loc * 2 + 1] <= low) return query(low, high, loc * 2 + 1);
		return max(query(low, r[loc * 2], loc * 2), query(l[loc * 2 + 1], high, loc * 2 + 1));
	}
	void init(ll x = 1) {
		if (x >= s) {
			l[x] = r[x] = x - s + 1;
			return;
		}
		init(x * 2);
		init(x * 2 + 1);
		l[x] = l[x * 2];
		r[x] = r[x * 2 + 1];
	}
	segtree(ll n) {
		N = n;
		s = 1LL << (ll)ceil(log2(N));
		tree.resize(2 * s + 1);
		l.resize(2 * s + 1);
		r.resize(2 * s + 1);
		init();
	}
};
vector<ll> adj[MAX], sav[MAX], rev[MAX];
ll dir[MAX], ddir[MAX], dddir[MAX];
vector<vector<ll>> chain;
vector<set<ll>> subtree, revtree;
vector<segtree> chainseg;
ll C[MAX], depth[MAX], mxdepv[MAX], prtval[MAX];
ll ord[MAX];
pll range[MAX];
ll sp[MAX][MAXS];
ll cnt;
ll num[MAX];
pll cnum[MAX];
ll ans[MAX];
ll init(ll x, ll p = 0, ll d = 0) {
	ord[x] = ++cnt;
	sp[x][0] = p;
	ll i;
	for (i = 1; i < MAXS; i++) sp[x][i] = sp[sp[x][i - 1]][i - 1];
	depth[x] = d;
	ll mx = 0;
	ll sum = 1;
	sav[x].resize(adj[x].size());
	rev[x].resize(adj[x].size());
	for (auto v : adj[x]) {
		if (v == p) continue;
		sum += init(v, x, d + 1);
	}
	range[x] = pll(ord[x], cnt);
	return num[x] = sum;
}
void calc(ll x, ll p = 0) {
	ll i;
	ll mx = 0;
	mxdepv[x] = x;
	for (i = 0; i < adj[x].size(); i++) {
		if (adj[x][i] == p) continue;
		calc(adj[x][i], x);
		if (mx < depth[mxdepv[adj[x][i]]]) mx = depth[mxdepv[adj[x][i]]], mxdepv[x] = mxdepv[adj[x][i]];
		sav[x][i] = mxdepv[adj[x][i]];
	}
	ll cnt = 0;
	ll vv, vvv;
	vv = vvv = -1;
	ll nv = -1;
	ll nmx = 0;
	for (auto v : adj[x]) {
		if (v == p) continue;
		if (depth[mxdepv[v]] == mx) cnt++, vvv = vv, vv = v;
		else if (depth[mxdepv[v]] > nmx) nmx = depth[mxdepv[v]], nv = v;
	}
	if (cnt >= 2) {
		for (i = 0; i < adj[x].size(); i++) {
			if (adj[x][i] != p) {
				rev[x][i] = (adj[x][i] == vv ? vvv : vv);
			}
		}
	}
	else {
		for (i = 0; i < adj[x].size(); i++) {
			if (adj[x][i] != p) {
				rev[x][i] = (adj[x][i] == vv ? nv : vv);
			}
		}
	}
	for (i = 0; i < adj[x].size(); i++) {
		if (adj[x][i] == p) continue;
		if (rev[x][i] == -1) rev[x][i] = x;
		else rev[x][i] = mxdepv[rev[x][i]];
	}
}
void make_chain(ll x, ll p = 0) {
	ll mx, mv;
	mx = mv = 0;
	chain[cnt].push_back(x);
	cnum[x] = { cnt, chain[cnt].size() - 1 };
	for (auto v : adj[x]) {
		if (v == p) continue;
		if (mx < num[v]) mx = num[v], mv = v;
	}
	if (mv) make_chain(mv, x);
	for (auto v : adj[x]) {
		if (v == p || v == mv) continue;
		cnt++;
		chain.push_back(vector<ll>());
		make_chain(v, x);
	}
}
void make_tree() {
	ll i;
	for (i = 0; i < chain.size(); i++) chainseg.push_back(segtree(chain[i].size()));
}
void update(ll v, ll x) {
	chainseg[cnum[v].first].update(cnum[v].second + 1, x);
}
//1을 루트로 하는 LCA
ll lca(ll u, ll v) {
	if (depth[u] != depth[v]) {
		if (depth[u] < depth[v]) swap(u, v);
		ll i;
		for (i = MAXS - 1; i >= 0; i--) if (depth[sp[u][i]] >= depth[v]) u = sp[u][i];
	}
	if (u == v) return u;
	ll i;
	for (i = MAXS - 1; i >= 0; i--) if (sp[u][i] != sp[v][i]) u = sp[u][i], v = sp[v][i];
	return sp[v][0];
}
//HLD query
ll mxval(ll u, ll v) {
	ll ans = 0;
	ll l = lca(u, v);
	while (cnum[u].first != cnum[l].first) ans = max(ans, chainseg[cnum[u].first].query(1, cnum[u].second + 1)), u = sp[chain[cnum[u].first][0]][0];
	while (cnum[v].first != cnum[l].first) ans = max(ans, chainseg[cnum[v].first].query(1, cnum[v].second + 1)), v = sp[chain[cnum[v].first][0]][0];
	ans = max(ans, chainseg[cnum[l].first].query(cnum[l].second + 1, cnum[u].second + 1));
	ans = max(ans, chainseg[cnum[l].first].query(cnum[l].second + 1, cnum[v].second + 1));
	return ans;
}
//두 정점 사이 거리
ll dis(ll u, ll v) { return depth[u] + depth[v] - 2 * depth[lca(u, v)]; }
//r이 루트, v의 x번째 부모
ll prtx(ll r, ll v, ll x) {
	if (x == 0) return v;
	ll rv = dis(r, v);
	if (rv < x) return 0;
	ll l = lca(r, v);
	if (dis(l, v) < x) {
		ll d = rv - x;
		ll i;
		for (i = MAXS - 1; i >= 0; i--) if (d - (1 << i) >= 0) d -= (1 << i), r = sp[r][i];
		return r;
	}
	else {
		ll i;
		for (i = MAXS - 1; i >= 0; i--) if (x - (1 << i) >= 0) x -= (1 << i), v = sp[v][i];
		return v;
	}
}
ll getfar(ll v, ll ban) {
	if (dir[v] != ban) return dir[v];
	return ddir[v];
}
ll getfar(ll v, ll ban1, ll ban2) {
	if (ban1 > ban2) swap(ban1, ban2);
	if (ban2 == -1) return dir[v];
	if (ban1 == -1) return getfar(v, ban2);
	if (dir[v] != ban1 && dir[v] != ban2) return dir[v];
	if (ddir[v] != ban1 && ddir[v] != ban2) return ddir[v];
	return dddir[v];
}
ll getind(vector<ll>& v, ll c) {
	return lower_bound(v.begin(), v.end(), c) - v.begin();
}
//r1 : previous root
void prop(ll r1, ll r2) {
	if (ddir[r2] == -1) update(r2, dis(sav[r2][dir[r2]], r2));
	else update(r2, dis(r2, sav[r2][dir[r2]]) + dis(r2, sav[r2][ddir[r2]]));
	ll ind = getind(adj[r1], r2);
	ll f1 = getfar(r1, ind);
	ll f2 = getfar(r1, ind, f1);
	if (f1 == -1) update(r1, 0);
	else if (f2 == -1) update(r1, dis(r1, sav[r1][f1]));
	else update(r1, dis(r1, sav[r1][f1]) + dis(r1, sav[r1][f2]));
}
void dfs(ll x, ll p = 0) {
	ll i;
	for (i = 0; i < adj[x].size(); i++) {
		ll fardir = getfar(adj[x][i], getind(adj[adj[x][i]], x));
		ll farv = sav[x][i];
		ll fardis = dis(x, farv);
		ll res = mxval(farv, adj[x][i]);
		if (res >= fardis) continue;
		ll xx = (fardis - 1) / 2;
		ll root = prtx(x, farv, xx);
		if (lca(root, x) == root) revtree[prtx(x, root, 1)].insert(C[x]);
		else subtree[root].insert(C[x]);
	}
	for (auto v : adj[x]) {
		if (v == p) continue;
		prop(x, v);
		dfs(v, x);
		prop(v, x);
	}
}
unordered_map<ll, ll> mp;
void getans(ll x, ll p = 0) {
	ll i;
	for (auto c : subtree[x]) mp[c]++;
	vector<ll> vv;
	for (auto c : revtree[x]) {
		mp[c]--;
		if (!mp[c]) vv.push_back(c);
	}
	for (auto x : vv) mp.erase(x);
	vv.clear();
	ans[x] = mp.size();
	for (auto v : adj[x]) {
		if (v == p) continue;
		getans(v, x);
	}
	for (auto c : revtree[x]) mp[c]++;
	for (auto c : subtree[x]) {
		mp[c]--;
		if (!mp[c]) vv.push_back(c);
	}
	for (auto x : vv) mp.erase(x);
}
void calcp(ll x, ll p = 0) {
	ll i, v;
	if (x != 1) {
		if (p == 1) sav[x][getind(adj[x], p)] = rev[p][getind(adj[p], x)];
		else {
			ll v1 = rev[p][getind(adj[p], x)];
			ll v2 = sav[p][getind(adj[p], sp[p][0])];
			if (v1 > v2) swap(v1, v2);
			if (dis(x, v1) >= dis(x, v2)) sav[x][getind(adj[x], p)] = v1;
			else sav[x][getind(adj[x], p)] = v2;
		}
	}
	for (auto v : adj[x]) if (v != p) calcp(v, x);
}
signed main() {
	ios::sync_with_stdio(false), cin.tie(0);
	depth[0] = -1;
	ll N, M;
	cin >> N >> M;
	ll i, j;
	ll a, b;
	for (i = 1; i < N; i++) cin >> a >> b, adj[a].push_back(b), adj[b].push_back(a);
	for (i = 1; i <= N; i++) cin >> C[i];
	for (i = 1; i <= N; i++) sort(adj[i].begin(), adj[i].end());
	init(1);
	calc(1);
	calcp(1);
	cnt = 0;
	chain.push_back(vector<ll>());
	make_chain(1);
	make_tree();
	//여기까지는 문제없음
	for (i = 1; i <= N; i++) {
		ll mx = 0, mv = 0;
		dir[i] = ddir[i] = dddir[i] = -1;
		for (j = 0; j < adj[i].size(); j++) if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dir[i] = j;
		mx = 0;
		for (j = 0; j < adj[i].size(); j++) {
			if (j == dir[i]) continue;
			if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), ddir[i] = j;
		}
		mx = 0;
		for (j = 0; j < adj[i].size(); j++) {
			if (j == dir[i] || j == ddir[i]) continue;
			if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dddir[i] = j;
		}
		if (i != 1) {
			ll p = getind(adj[i], sp[i][0]);
			ll r = getfar(i, p);
			ll rr = getfar(i, p, r);
			ll xx = 0;
			if (r != -1) xx += dis(i, sav[i][r]);
			if (rr != -1) xx += dis(i, sav[i][rr]);
			update(i, xx);
		}
		else {
			ll dd;
			dd = ddir[1];
			if (dd == -1) update(1, depth[sav[1][dir[1]]]);
			else update(1, depth[sav[1][dd]] + depth[sav[1][dir[1]]]);
		}
	}
	subtree.resize(N + 1);
	revtree.resize(N + 1);
	dfs(1);
	for (i = 1; i <= N; i++) for (auto c : revtree[i]) mp[c]++;
	getans(1);
	for (i = 1; i <= N; i++) cout << ans[i] << ln;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'll init(ll, ll, ll)':
joi2019_ho_t5.cpp:69:5: warning: unused variable 'mx' [-Wunused-variable]
   69 |  ll mx = 0;
      |     ^~
joi2019_ho_t5.cpp: In function 'void calc(ll, ll)':
joi2019_ho_t5.cpp:84:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   84 |  for (i = 0; i < adj[x].size(); i++) {
      |              ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:101:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  101 |   for (i = 0; i < adj[x].size(); i++) {
      |               ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:108:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  108 |   for (i = 0; i < adj[x].size(); i++) {
      |               ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:114:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  114 |  for (i = 0; i < adj[x].size(); i++) {
      |              ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'void make_tree()':
joi2019_ho_t5.cpp:139:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::vector<long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  139 |  for (i = 0; i < chain.size(); i++) chainseg.push_back(segtree(chain[i].size()));
      |              ~~^~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'void dfs(ll, ll)':
joi2019_ho_t5.cpp:214:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  214 |  for (i = 0; i < adj[x].size(); i++) {
      |              ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:215:6: warning: unused variable 'fardir' [-Wunused-variable]
  215 |   ll fardir = getfar(adj[x][i], getind(adj[adj[x][i]], x));
      |      ^~~~~~
joi2019_ho_t5.cpp: In function 'void getans(ll, ll)':
joi2019_ho_t5.cpp:234:5: warning: unused variable 'i' [-Wunused-variable]
  234 |  ll i;
      |     ^
joi2019_ho_t5.cpp: In function 'void calcp(ll, ll)':
joi2019_ho_t5.cpp:256:5: warning: unused variable 'i' [-Wunused-variable]
  256 |  ll i, v;
      |     ^
joi2019_ho_t5.cpp:256:8: warning: unused variable 'v' [-Wunused-variable]
  256 |  ll i, v;
      |        ^
joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:290:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  290 |   for (j = 0; j < adj[i].size(); j++) if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dir[i] = j;
      |               ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:292:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  292 |   for (j = 0; j < adj[i].size(); j++) {
      |               ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:297:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  297 |   for (j = 0; j < adj[i].size(); j++) {
      |               ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:288:14: warning: unused variable 'mv' [-Wunused-variable]
  288 |   ll mx = 0, mv = 0;
      |              ^~

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