版权声明:LeoJAM Presents https://blog.csdn.net/fcb_x/article/details/82807735
给定一颗n个结点的无根树,树上的每个点有一个非负整数点权,定义一条路径的价值为路径上的点权和-路径的点权最大值。
给定参数p,我们想知道,有多少不同的树上简单路径,满足它的价值恰好是p的倍数。
注意:单点算作一个路径;u ≠ v时,(u,v)和(v,u)只算一次。
这个实际上是聪聪可可的升级版本
它和一般的点分治不太一样
他的代价是在点权上的。
这个存在一个问题:
你维护两条点权路径,会出现LCA的点权被计算了两次所以注意减去点权统计
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
typedef int INT;
#define int long long
using namespace std;
const int N=2e6+100;
struct Front_star{
int u,v,nxt;
}e[N<<1];
int cnt=0;
int first[N];
void add(int u,int v){
++cnt;
e[cnt].u=u;
e[cnt].v=v;
e[cnt].nxt=first[u];
first[u]=cnt;
}
//
int n,P;
int All;
int root;
int ans;
int siz[N];
int vis[N];
int val[N];
int F[N];
void Get_Root(int u,int fat){
int now=0;
siz[u]=1;
for(int i=first[u];i;i=e[i].nxt){
int v=e[i].v;
if(v==fat||vis[v])continue;
Get_Root(v,u);
siz[u]+=siz[v];
if(now<siz[v])now=siz[v];
}
if(now<All-siz[u])now=All-siz[u];
F[u]=now;
if(F[root]>now)root=u;
}
int Mx[N];
int d[N];
int com[N];
void Get_Dep(int u,int fat,int dep,int M){
for(int i=first[u];i;i=e[i].nxt){
int v=e[i].v;
if(vis[v]||v==fat)continue;
Mx[++com[0]]=max(M,val[v]);
d[com[0]]=(dep+val[v])%P;
Get_Dep(v,u,d[com[0]],Mx[com[0]]);
}
}
bool cmp(int A,int B){
return Mx[A]<Mx[B];
}
int bin[N*10];
int Solve(int u,int dep,int M,int v){
// cout<<u<<" "<<dep<<" "<<M<<'\n';
d[com[0]=1]=(val[u]+dep)%P;
Mx[com[0]]=max(M,val[u]);
Get_Dep(u,0,d[com[0]],Mx[com[0]]);
int ret=0;
for(int i=1;i<=com[0];++i){
com[i]=i;
}
sort(com+1,com+1+com[0],cmp);
// for(int i=1;i<=com[0];++i){
// cout<<Mx[i]<<" "<<d[i]<<'\n';
// }
for(int i=1;i<=com[0];++i){
ret+=bin[(((P-(d[com[i]]-Mx[com[i]])%P+P)%P)+P)%P];
bin[((d[com[i]]-v)%P+P)%P]++;;
}
for(int i=1;i<=com[0];++i){
bin[((d[com[i]]-v)%P+P)%P]=0;
}
// cout<<ret<<'\n';
return ret;
}
void DFS(int u){
// cout<<u<<'\n';
vis[u]=1;
ans+=Solve(u,0,0,val[u]);
// exit(0);
for(int i=first[u];i;i=e[i].nxt){
int v=e[i].v;
if(vis[v])continue;
ans-=Solve(v,val[u],val[u],val[u]);
All=siz[v];
root=0;
Get_Root(v,u);
DFS(root);
}
}
INT main(){
//freopen("path.in","r",stdin);
//freopen("path.out","w",stdout);
F[0]=1e9+7;
scanf("%lld%lld",&n,&P);
for(int i=1;i<n;++i){
int u,v;
scanf("%lld%lld",&u,&v);
add(u,v);
add(v,u);
}
for(int i=1;i<=n;++i)scanf("%lld",&val[i]);
All=n;
root=0;
Get_Root(1,0);
// cout<<root<<'\n';
DFS(root);
cout<<ans+n;
return 0;
}