洛谷P3806 【模板】点分治1

题链

点分治分为四步：
1.找到树的重心
2.删除树的重心
3.处理经过重心的路径
4.处理重心的子树

详解来自BiliBili

#include<bits/stdc++.h>
using namespace std;
#define LL long long 
#define MAXN 10000009
#define MS 1000009

int n,m;
struct node{
	int to;
	LL val;
};
vector<node > vc[MS];
int ask[MS] ,ac[MS];
int rt,tr_size;
int sz[MS],w[MS];
int del[MS];
int dis[MS];
int dislist[MS],cntd;
int ext[MAXN];
queue<int > tmp;

void get_rt(int u,int f){ // 找重心 
	sz[u] = 1;
	w[u] = 0;
	for(auto &nb:vc[u]){
		int v = nb.to;
		if(v != f && !del[v]){
			get_rt(v,u);
			sz[u] += sz[v];
			w[u] = max(w[u],sz[v]); // 得到最大的子树大小 
		}
	}
	w[u] = max(w[u] ,tr_size-sz[u]); // 自己以及祖先的大小 
	if(w[u] < w[rt]) rt = u; // 更新重心 
}

void get_dis(int u,int f){
	dislist[++cntd] = dis[u]; // 记录所有距离 
	for(auto &nb:vc[u]){
		int v = nb.to;
		LL val = nb.val;
		if(v != f && !del[v]){
			dis[v] = dis[u] + val;
			get_dis(v,u);
		}
	}
}

void cal(int u){
	ext[0] = 1; // 标记权值 0 存在 
	for(auto &nb:vc[u]){
		int v = nb.to;
		LL val = nb.val;
		if(!del[v]){
			// 得到 v为根的子树到 u 的所有距离 
			cntd = 0;
			dis[v] = val;
			get_dis(v,u);
			// 处理询问 
			for(int i=1;i<=m;i++){
				for(int j=1;j<=cntd;j++){
					if(ask[i] >= dislist[j]){
						ac[i] |= ext[ ask[i] - dislist[j] ];
					}
				}
			}
			// 记录所有距离，标记该距离存在 
			for(int i=1;i<=cntd;i++){
				if(dislist[i] <= 1e7){
					ext[ dislist[i] ] = 1;
					tmp.push( dislist[i] );
				}
			}
		}
	}
	// 删除所有距离 
	ext[0] = 0;
	while(!tmp.empty()){
		ext[ tmp.front() ] = 0;
		tmp.pop();
	}
}

void divide(int u){
	del[u] = 1; // 删除重心 
	cal(u);		// 处理经过重心的路径 
	for(auto &nb:vc[u]){ // 对于每一个子树同样分治 
		int v = nb.to;
		if(!del[v]){
			// 找重心并分治 
			w[rt = 0] = tr_size = sz[v];
			get_rt(v,0);
			get_rt(rt,0);
			divide(rt);
		}
	}
}

void solve(){
	cin >> n >> m;
	for(int i=1;i<=n-1;i++){
		int u,v;
		LL val;
		cin >> u >> v >> val;
		vc[u].push_back({v,val});
		vc[v].push_back({u,val});  
	}
	for(int i=1;i<=m;i++){
		cin >> ask[i];
	}
	// 找到树的重心 
	w[rt = 0] = tr_size = n;
	get_rt(1,0);
	get_rt(rt,0);
	// 分治 
	divide(rt);
	
	for(int i=1;i<=m;i++){
		if(ac[i]) cout << "AYE\n";
		else cout << "NAY\n";
	}
}

int main(){
	ios::sync_with_stdio(false);
	int ce;
//	cin >> ce;
	ce = 1;
	while(ce--){
		solve();
	}
	
	return 0;
}