点分治学习笔记
呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕呕
Part 0 树的重心
定义：以某一个点作为根节点，若这个点所有子树大小的最大值最小，则称这个点为树的重心。
性质
1.树的重心如果不唯一，则之多有两个，且这两个重心相邻。可以类比中位数理解。
2.以树的重心为根时，所有子树的大小都不超过整棵树大小的一般。可以反证。
3.树中所有点到某个点的距离和中，到中心的距离和是最小的（显然）。如果有两个重心，那么到它们的距离和一样。
4.把两棵树通过一条边相连得到一颗新的树，那么新的树的重心在连接原来两颗树的重心的路径上。
5.在一棵树上添加或删除一个叶子，那么它的重心最多只移动一条边的距离。
例题
比较老的题，就是让你求每个子树的重心。我们可以用上面的性质2，4。考虑从底向上递归，假设我们已经知道了前几个子树的重心 \(p\)，合并之后的树的重心肯定在新来的子树的重心和 \(p\) 之间的连边上，那我们就只需要暴力跳，用性质 2 判断就行了。不要小看这个题，后边点分树还要用到这个处理过程。
#include<iostream>
#include<cstdio>
using namespace std;
const int N=300005;
int n,q,h[N],tot,fa[N],siz[N],hson[N],ans[N];
struct edge{int v,nxt;}e[N<<1];
void add(int u,int v){e[++tot].v=v,e[tot].nxt=h[u],h[u]=tot;}
void dfs(int u,int fath)
{
ans[u]=u,siz[u]=1,hson[u]=0;
for(int i=h[u],v;i;i=e[i].nxt)if((v=e[i].v)!=fath)
{
dfs(v,u);siz[u]+=siz[v];
if(!hson[u]||siz[v]>siz[hson[u]])hson[u]=v;
}
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fath)continue;
for(int p=ans[v];p!=fa[u];p=fa[p])
{
if(max(siz[hson[p]],siz[u]-siz[p])<=siz[u]/2)
{
ans[u]=p;break;
}
}
}
return;
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0);cout.tie(0);
cin>>n>>q;
for(int i=2;i<=n;i++){cin>>fa[i];add(fa[i],i);}
dfs(1,0);
while(q--){int x;cin>>x;cout<<ans[x]<<"\n";}
return 0;
}
Part 1 点分治
来看板子。
树的点分治的基本思想为：
1.找到树的一个重心，则问题的解可以分为路径经过重心的解和路径不经过中心的解两部分。
2.路径经过重心的解：设重心为树根，用数据结构记录重心到每个节点的路径信息。从中取出根的不同子树中的路径信息组合起来即为路径经过重心的解。
3.路径不经过重心的解：将重心和其关联的边删除，将分割出来的每棵树递归处理即可求出路径不经过重心的解.
#include<iostream>
#include<cstdio>
using namespace std;
const int N=1e4+10,M=1e7+5;
const int INF=0x3f3f3f3f;
int n,m,que[N];
int h[N],tot;
struct edge{
int v,w,nxt;
}e[N<<1];
void add(int u,int v,int w)
{
e[++tot].v=v,e[tot].w=w;
e[tot].nxt=h[u],h[u]=tot;
}
int siz[N],dis[N];
bool res[N];
bool vis[N];
int sum,Temp,rt;
void find_rt(int u,int fa)
{
siz[u]=1;//一定不要忘了写这一行
int mx=0;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fa||vis[v])continue;
find_rt(v,u);
siz[u]+=siz[v];
if(siz[v]>mx)mx=siz[v];
}
mx=max(mx,sum-siz[u]);//注意这里是用 sum 减，不是 n
if(mx<Temp)Temp=mx,rt=u;
}
int rem[N],rtot;
void calc_dis(int u,int fa)
{
rem[++rtot]=dis[u];
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fa||vis[v])continue;
dis[v]=dis[u]+e[i].w;
calc_dis(v,u);
}
}
bool chk[10000005];
int q[N],p;
void calc(int u)
{
p=0;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
rtot=0,dis[v]=e[i].w;
calc_dis(v,u);
//下面两个循环顺序不要反了，否则结果是错的
for(int j=1;j<=rtot;j++)
for(int k=1;k<=m;k++)
if(que[k]>=rem[j])
res[k]|=chk[que[k]-rem[j]];
for(int j=1;j<=rtot;j++)
if(rem[j]<=1e7)
q[++p]=rem[j],chk[rem[j]]=1;
}
for(int i=1;i<=p;i++)chk[q[i]]=0;
}
void solve(int u)
{
vis[u]=chk[0]=1;
calc(u);
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
sum=siz[v],Temp=INF;
find_rt(v,0);
solve(rt);//不要写成 solve(v)
}
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0);cout.tie(0);
cin>>n>>m;
for(int i=1;i<n;i++)
{
int u,v,w;
cin>>u>>v>>w;
add(u,v,w);
add(v,u,w);
}
for(int i=1;i<=m;i++)
{
cin>>que[i];
if(que[i]==0)res[i]=1;
}
sum=n,Temp=INF;//不要忘了
find_rt(1,0);
solve(rt);
for(int i=1;i<=m;i++)
{
if(res[i])cout<<"AYE\n";
else cout<<"NAY\n";
}
return 0;
}
放一个压了行的板子。
#include<iostream>
using namespace std;
int n,m,h[10005],tot,k[105];
struct edge{int v,w,nxt;}e[20005];
void add(int u,int v,int w){e[++tot].v=v,e[tot].w=w,e[tot].nxt=h[u],h[u]=tot;}
struct DFZ{
bool vis[10005],ans[105],hv[10000005];
int sum,rt,tmp,siz[10005],dis[10005],rem[100005],rtot,q[100005],p;
void find_rt(int u,int fath){
int mx=0;siz[u]=1;
for(int i=h[u],v;i;i=e[i].nxt){
if((v=e[i].v)==fath||vis[v])continue;
find_rt(v,u);siz[u]+=siz[v],mx=max(mx,siz[v]);
}if(max(mx,sum-siz[u])<tmp)tmp=max(mx,sum-siz[u]),rt=u;
}void calc_dis(int u,int fath){
rem[++rtot]=dis[u];
for(int i=h[u],v;i;i=e[i].nxt){
if((v=e[i].v)==fath||vis[v])continue;
dis[v]=dis[u]+e[i].w;calc_dis(v,u);
}
}void calc(int u){p=0;
for(int i=h[u],v;i;i=e[i].nxt){
if(vis[(v=e[i].v)])continue;
dis[v]=e[i].w,rtot=0;calc_dis(v,u);
for(int j=1;j<=rtot;j++)for(int l=1;l<=m;l++)if(k[l]>=rem[j])ans[l]|=hv[k[l]-rem[j]];
for(int j=1;j<=rtot;j++)if(rem[j]<=1e7)hv[rem[j]]=true,q[++p]=rem[j];
}for(int i=1;i<=p;i++)hv[q[i]]=false;
}void solve(int u){
vis[u]=hv[0]=true;calc(u);
for(int i=h[u],v;i;i=e[i].nxt){
if(vis[(v=e[i].v)])continue;
sum=siz[v],tmp=0x3f3f3f3f;
find_rt(v,0);solve(rt);
}
}void mian(){sum=n,tmp=0x3f3f3f3f;find_rt(1,0);solve(rt);for(int i=1;i<=m;i++)puts(((ans[i]||!k[i])?"AYE":"NAY"));}
}dfz;
int main(){
scanf("%d%d",&n,&m);
for(int i=1;i<n;i++){int u,v,w;scanf("%d%d%d",&u,&v,&w);add(u,v,w);add(v,u,w);}
for(int i=1;i<=m;i++)scanf("%d",k+i); dfz.mian();
return 0;
}
下面是一些习题
Race
改一下统计答案就行了。
#include<iostream>
#include<cstring>
#include<cstdio>
using namespace std;
const int N=2e5+10,M=1e6+10;
const int INF=0x3f3f3f3f;
int n,k,h[N],tot,ans=INF;
bool flag=false;
struct edge{
int v,w,nxt;
}e[N<<1];
void add(int u,int v,int w)
{
e[++tot].v=v,e[tot].w=w;
e[tot].nxt=h[u],h[u]=tot;
}
int rt,tmp=INF,siz[N],dis[N],sum=n;
bool vis[N];
void get_rt(int u,int fath)
{
siz[u]=1;int mx=0;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fath||vis[v])continue;
get_rt(v,u);
siz[u]+=siz[v];
mx=max(mx,siz[v]);
}
mx=max(mx,sum-siz[u]);
if(mx<tmp)tmp=mx,rt=u;
}
int rtot;
int rem[N],cc[N],f[M];
void calc_dis(int u,int fath,int dep)
{
if(dis[u]<=1e6)rem[++rtot]=dis[u],cc[rtot]=dep;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fath||vis[v])continue;
dis[v]=dis[u]+e[i].w;
calc_dis(v,u,dep+1);
}
}
int q[N],qtot;
void calc(int u)
{
qtot=0;f[0]=0;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
dis[v]=e[i].w;
rtot=0;
calc_dis(v,u,1);
for(int j=1;j<=rtot;j++)if(k>=rem[j])
ans=min(ans,f[k-rem[j]]+cc[j]);
for(int j=1;j<=rtot;j++)q[++qtot]=rem[j],f[rem[j]]=min(f[rem[j]],cc[j]);
}
for(int i=1;i<=qtot;i++)f[q[i]]=INF;
}
void solve(int u)
{
vis[u]=1;
calc(u);
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
sum=siz[v],tmp=INF;
get_rt(v,0);
solve(rt);
}
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0);cout.tie(0);
memset(f,INF,sizeof(f));
cin>>n>>k;
for(int i=1;i<n;i++)
{
int u,v,w;
cin>>u>>v>>w;
u++,v++;
add(u,v,w);
add(v,u,w);
}
sum=n,tmp=INF;
get_rt(1,0);
solve(rt);
if(ans==INF)cout<<-1<<"\n";
else cout<<ans<<"\n";
return 0;
}
Tree
统计答案的时候开一个树状数组，能把时间复杂度降低，从而过掉这个题。
这个题可以启发我们在后面的题里写点分治套线段树这种恶心东西。
代码：
#include<cstdio>
#define N 40005
#define max(a,b) (a)>=(b)?(a):(b)
#define min(a,b) (a)<=(b)?(a):(b)
using namespace std;
int n,k,h[N],tot;
struct edge{
int v,w,nxt;
}e[N<<1];
void add(int u,int v,int w)
{
e[++tot].v=v,e[tot].w=w;
e[tot].nxt=h[u],h[u]=tot;
}
struct BIT{
int c[N],m;
BIT(){for(int i=0;i<N;i++)c[i]=0;}
inline int lowbit(int x){return x & -x;}
void clear(){for(int i=0;i<=m;i++)c[i]=0;}
void upd(int x,int k){while(x<=m)c[x]+=k,x+=lowbit(x);}
int ask(int x){int res=0;while(x)res+=c[x],x-=lowbit(x);return res;}
}tr;
bool vis[N];
int siz[N],tem,rt,sum;
void find_rt(int u,int fath)
{
int mx=0;siz[u]=1;//最开始就是因为没写 siz[u]=1 TLE了一页，再次敲警钟
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fath||vis[v])continue;
find_rt(v,u);siz[u]+=siz[v];
mx=max(mx,siz[v]);
}
mx=max(mx,sum-siz[u]);
if(mx<tem)tem=mx,rt=u;
}
int dis[N],rem[N],rtot,ans;
void calc_dis(int u,int fath)
{
if(dis[u]<=k)rem[++rtot]=dis[u];
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==fath||vis[v])continue;
dis[v]=dis[u]+e[i].w;
calc_dis(v,u);
}
}
void calc(int u)
{
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
dis[v]=e[i].w,rtot=0;
calc_dis(v,u);
for(int j=1;j<=rtot;j++)ans+=tr.ask(k-rem[j]);
for(int j=1;j<=rtot;j++){tr.upd(rem[j],1);ans++;}
}
tr.clear();
}
void solve(int u)
{
calc(u);vis[u]=true;
for(int i=h[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
sum=siz[v],tem=0x3f3f3f3f;
find_rt(v,u);solve(rt);
}
}
int main()
{
scanf("%d",&n);
for(int i=1;i<n;i++)
{
int u,v,w;
scanf("%d%d%d",&u,&v,&w);
add(u,v,w);add(v,u,w);
}
scanf("%d",&k);tr.m=k;
sum=n,tem=0x3f3f3f3f;
find_rt(1,0);solve(rt);
printf("%d\n",ans);
return 0;
}
Transport
这道题中我们分类别来讨论路径。设当前的根是 \(rt\)，点 \(u\) 到 \(rt\) 的点权和为 \(w_u\),边权和是 \(d_u\)，父节点为 \(fa_u\)。
第一种情况是 \(u \rightarrow rt\)，考虑该路径合法的条件是对于任意路径上的点 \(v\) 都有 \(d_u-d_v \leq w_u-w_v\)，换句话说就是 \(u\) 到路径上任意一点时油量都大于等于 0。将上面的式子变形，得到 \(w_v-d_v \leq w_u-d_u\)，那我们只需要让路径上最大的 \(w_v-d_v\) 小于 \(w_u-d_u\) 即可满足条件。于是我们只需记 \(w_u-d_u\) 的最大值就可以判断点 \(u\) 能不能走到 \(rt\)。
第二种情况 \(rt \rightarrow u\) 就有点儿复杂，因为这时货车的起点不一定是 \(rt\)，如果总路径是 \(v \rightarrow rt \rightarrow u\) ，设 \(v \rightarrow rt\) 路径结束后剩余油量为 \(x\)，则路径 \(rt \rightarrow v\) 合法的条件是对于路径上任意一点 \(k\)，都有 \(x+w_{fa_k}-d_k \geq 0\)。这里我们可以记 \(y=min \{w_{fa_u}-d_u\}\)，则该路径合法的充分必要条件就是 \(x+y \geq 0\)。
开两个队列q1, q2,其中 q1 记所有合法路径到 \(rt\) 时的剩余油量，q2 记所有下行路径的 \(min\{w_{fa_u}-d_u\}\)，将两个队列中的元素升序排列，倒序枚举 q2,看 q1 中多少值满足下行路径的合法条件。这个可以用双指针做到 \(O(N)\) 复杂度。当然，如果我们用这种方法合并路径，会将同一子树内的条件也算上，这需要我们把他们排掉。具体可以看代码。
有一个细节就是我们在合并两条路径的时候，因为这两条路径都算上了 \(a_{rt}\)，所以我们要利用容斥原理减掉一个 \(a_{rt}\)。还有就是点分治本来复杂度就不算低，合并路径过程中复杂度当然是越低越保险，所以这里并不推荐用二分的方法平添一个 \(\log\)。
代码：
#include<iostream>
#include<cstdio>
#include<algorithm>
#define int long long
using namespace std;
const int N=200005;
int n,a[N],ans;
int head[N],tot;
struct edge{
int v,w,nxt;
}e[N<<1];
void add(int u,int v,int w)
{
e[++tot].v=v,e[tot].w=w;
e[tot].nxt=head[u],head[u]=tot;
}
bool vis[N];
int siz[N],rt,tem,sum;
void find_rt(int u,int fath)
{
int mx=0;
siz[u]=1;
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v]||v==fath)continue;
find_rt(v,u);
siz[u]+=siz[v];
mx=max(mx,siz[v]);
}
mx=max(mx,sum-siz[u]);
if(mx<tem)tem=mx,rt=u;
}
int w[N],d[N];
int up[N],down[N],cnt1,cnt2;
void calc_dis(int u,int fath,int mx,int mn)
{
if(w[u]-d[u]>=mx)up[++cnt1]=w[u]-d[u];
down[++cnt2]=mn;
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v]||v==fath)continue;
d[v]=d[u]+e[i].w,w[v]=w[u]+a[v];
calc_dis(v,u,max(mx,w[u]-d[u]),min(mn,w[u]-d[v]));
}
}
int q1[N],q2[N],p1,p2;
inline void calc(int u)
{
p1=p2=0;
d[u]=0,w[u]=a[u];
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
cnt1=cnt2=0;
w[v]=w[u]+a[v],d[v]=e[i].w;
calc_dis(v,u,w[u]-d[u],w[u]-d[v]);
sort(up+1,up+1+cnt1);
sort(down+1,down+1+cnt2);
int l=1;
for(int i=cnt2;i>=1;i--)
{
while(l<=cnt1&&up[l]+down[i]-a[u]<0)l++;
ans-=cnt1-l+1;
}
for(int i=1;i<=cnt1;i++)q1[++p1]=up[i];
for(int i=1;i<=cnt2;i++)q2[++p2]=down[i];
}
sort(q1+1,q1+1+p1);
sort(q2+1,q2+1+p2);
int l=1;
for(int i=p2;i>=1;i--)
{
if(q2[i]>=0)ans++;
while(l<=p1&&q1[l]+q2[i]-a[u]<0)l++;
ans+=p1-l+1;
}
ans+=p1;
}
void solve(int u)
{
vis[u]=true;
calc(u);
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(vis[v])continue;
sum=siz[v],tem=1000000000;
find_rt(v,u);
solve(rt);
}
}
signed main()
{
ios::sync_with_stdio(false);
cin.tie(0);cout.tie(0);
cin>>n;
for(int i=1;i<=n;i++)cin>>a[i];
for(int i=1;i<n;i++)
{
int u,v,w;cin>>u>>v>>w;
add(u,v,w);add(v,u,w);
}
sum=n,tem=1000000000;
find_rt(1,0);
solve(rt);
cout<<ans<<"\n";
return 0;
}
树的难题
考虑在统计答案的过程中，如果新来的的与待匹配的两者到根节点的颜色相等，那么结果就要减一。我们按颜色种类排序，这样相同颜色的就可以在一起了，开两个线段树分别统计与新来的边颜色相同与不同的结果最大值，找就行了，时间复杂度 \(O(Nlog^2N)\)（也可以用单调队列，时间复杂度降成 \(O(NlogN)\)）。
Part 2 点分树
本质上是对点分治的优化。我们把点分治中每个树的重心建立父子关系，重新构成的树就是点分树，这样我们就可以优化找根的过程。
