[TIOJ] 1696. Problem F 橘子園保衛戰

題目連結:http://tioj.infor.org/problems/1696

---

看來我對重心剖分的理解還不夠QQ，一直想成重心樹的父子關係會跟原本的樹一樣。
我對於每個重心樹上的點維護他重心子樹們中$ dis \leq i, \forall 0 \leq i \leq D_{max} $的人有幾個(其中dis代表在原樹上的距離)，以及扣掉某個小孩的子樹後的這坨值。接著對於每個人就當做一個詢問，從葉子往上走訪重心樹，同時查詢不走剛剛來的那個點的子樹，距離好的點有幾個。

#include <bits/stdc++.h>
using namespace std;
#define ALL(x) begin(x), end(x)
#define PB push_back
#define SZ(x) ((int)(x).size())
const int N = 100000 + 5;

struct node{
	int cur, dis;
	node(int a=-1,int b=0):cur(a),dis(b){}
};
vector<node> path[N];
vector<int> sum[N], fa_sum[N];
bool done[N];
int sz[N], M[N], fa[N], dis[N], que[N], cen[N];
vector<int> G[N];

int CenDe(int);
int Query(int,int);

int main(){
	ios_base::sync_with_stdio(0);cin.tie(0);
	int n; cin >> n;
	for(int i=1;i<=n;i++) cin >> que[i];
	for(int i=0;i<n-1;i++){
		int u, v; cin >> u >> v;
		G[u].PB(v); G[v].PB(u);
	}
	int cent = CenDe(1);
	partial_sum(ALL(sum[cent]), sum[cent].begin());
	for(int i=1;i<=n;i++) cout << Query(i, que[i]) << '\n';
	return 0;
}

int CenDe(int x){
	/* Start FINDING CENTROID */
	vector<int> bfs;
	{
		size_t pt = 0;
		bfs.PB(x); fa[x] = -1;
		while(pt < bfs.size()){
			int u = bfs[pt++];
			M[u] = 0; sz[u] = 1;
			for(int v: G[u]) if(!done[v] and fa[u] != v){
				fa[v] = u; bfs.PB(v);
			}
		}
	}
	reverse(ALL(bfs));
	for(int u: bfs) if(fa[u] != -1) {
		sz[fa[u]] += sz[u];
		M[fa[u]] = max(M[fa[u]], sz[u]);
	}
	for(int u: bfs) M[u] = max(M[u], sz[x] - sz[u]);
	int cent = -1;
	for(int u: bfs) if(M[u]*2 <= sz[x]) cent = u;
	/* End FINDING CENTROID */
	/* Start DECOMPOSITION */
	done[cent] = true;
	for(int u: G[cent]) if(!done[u]){
		cen[u] = CenDe(u);
	}
	done[cent] = false;
	/* End DECOMPOSITION */
	/* Start BUILD CENTROID TREE */
	for(int u: bfs) path[u].emplace_back(cent);
	bfs.clear();
	{
		size_t pt = 0;
		bfs.PB(cent); fa[cent] = -1; dis[cent] = 0;
		while(pt < bfs.size()){
			int u = bfs[pt++];
			if(SZ(sum[cent]) < dis[u]+1) sum[cent].resize(dis[u]+1);
			sum[cent][dis[u]] += 1;
			for(int v: G[u]) if(!done[v] and fa[u] != v){
				fa[v] = u; bfs.PB(v);
				dis[v] = dis[u]+1;
			}
			path[u].back().dis = dis[u];
		}
	}
	for(int u: G[cent]) if(!done[u]) {
		fa_sum[cen[u]] = sum[cent];
		done[cent] = true; bfs.clear();
		dis[u] = 1; fa[u] = -1;
		size_t pt = 0; bfs.PB(u);
		while(pt < bfs.size()) {
			int _ = bfs[pt++]; fa_sum[cen[u]][dis[_]] -= 1;
			for(int v: G[_]) if(!done[v] and fa[_] != v){
				fa[v] = _; bfs.PB(v);
				dis[v] = dis[_]+1;
			}
		}
		done[cent] = false;
		partial_sum(ALL(sum[cen[u]]), sum[cen[u]].begin());
		partial_sum(ALL(fa_sum[cen[u]]), fa_sum[cen[u]].begin());
	}
	/* End BUILD CENTROID TREE */
	return cent;
}

int Query(int id, int len){
	int ret = 0, prev = -1;
	for(auto e: path[id]) {
		int x = e.cur, d = e.dis;
		if(prev == -1 and d<=len) ret += sum[x][min(SZ(sum[x])-1, len-d)];
		else if(d<=len) ret += fa_sum[prev][min(SZ(fa_sum[prev])-1, len-d)];
		prev = x;
	}
	return ret;
}