BZOJ 1316: 树上的询问 （点分治+set）

题目：http://www.lydsy.com/JudgeOnline/problem.php?id=1316

因为只要求存在某条路径长度为K，所以点分，然后用set判断差值是否在set中就可以了。

#include<cstring>
#include<algorithm>
#include<iostream>
#include<cstdio>
#include<queue>
#include<set>
#include<cmath>
#define rep(i,l,r) for (int i=l;i<=r;i++)
#define down(i,l,r) for (int i=l;i>=r;i--)
#define clr(x,y) memset(x,y,sizeof(x))
#define maxn 10900
#define inf int(1e9)
#define mm 1000000007
#define esp 1e-6
using namespace std;
#define ll long long
struct data{int obj,pre; ll c;
}e[maxn*2];
int head[maxn],s[maxn],ans[maxn],q[maxn],vis[maxn];
int n,m,tot,sum,mn,rt;
set<ll> bst;
int read(){
int x=0,f=1; char ch=getchar();
while (!isdigit(ch)) {if (ch=='-') f=-1; ch=getchar();}
while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();
return x*f;
}
void insert(int x,int y,int z){
e[++tot].obj=y; e[tot].pre=head[x]; e[tot].c=z; head[x]=tot;
}
void dfs(int u,int fa){
s[u]=1;
int mx=0;
for (int j=head[u];j;j=e[j].pre){
int v=e[j].obj;
if (!vis[v]&&v!=fa){
dfs(v,u);
s[u]+=s[v];
mx=max(mx,s[v]);
}
}
mx=max(sum-mx,mx);
if (mx<mn) mn=mx,rt=u;
}
void add(int u,int fa,ll w){
bst.insert(w);
for (int j=head[u];j;j=e[j].pre){
int v=e[j].obj;
if (v!=fa&&!vis[v])add(v,u,w+e[j].c);
}
}
void update(int u,int fa,ll w){
rep(i,1,m) if (!ans[i])
ans[i]|=(bst.find(q[i]-w)!=bst.end());
for (int j=head[u];j;j=e[j].pre){
int v=e[j].obj;
if (v!=fa&&!vis[v]) update(v,u,w+e[j].c);
}
}
void solve(int u){
mn=inf;
dfs(u,0);
u=rt;

bst.clear(); bst.insert(0);
for (int j=head[u];j;j=e[j].pre){
int v=e[j].obj;
if (!vis[v]) update(v,u,e[j].c),add(v,u,e[j].c); //先判断点再加边。
}

vis[u]=1;
for (int j=head[u];j;j=e[j].pre){
int v=e[j].obj;
if (!vis[v]){
sum=s[v];
solve(v);
}
}

}
int main(){
n=read(); m=read();
int x,y,z;
rep(i,1,n-1){
x=read(); y=read(); z=read();
insert(x,y,z);
insert(y,x,z);
}
rep(i,1,m) q[i]=read();
sum=n;

solve(1);
rep(i,1,m) if(!q[i]) ans[i]=1; //如果q[i]=0，那么可以只选取一个点。
rep(i,1,m) if (ans[i]) puts("Yes"); else puts("No");
return 0;
}