题目描述
给定一颗n个结点的无根树,树上的每个点有一个非负整数点权,定义一条路径的价值为路径上的点权和-路径的点权最大值。
给定参数p,我们想知道,有多少不同的树上简单路径,满足它的价值恰好是p的倍数。
注意:单点算作一个路径;u ≠ v时,(u,v)和(v,u)只算一次。
题解
随便点分治
同一个分治中心按到根最大值排序,然后维护桶
裸题
#include<cstdio>
#include<algorithm>
#define fo(i,a,b) for(i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int maxn=100000+10;
int a[maxn],mx[maxn],d[maxn],cnt[10000000+10],size[maxn],v[maxn];
int h[maxn],go[maxn*2],next[maxn*2];
bool bz[maxn];
int i,j,k,l,t,n,m,p,tot,top;
ll ans;
int read(){
int x=0,f=1;
char ch=getchar();
while (ch<'0'||ch>'9'){
if (ch=='-') f=-1;
ch=getchar();
}
while (ch>='0'&&ch<='9'){
x=x*10+ch-'0';
ch=getchar();
}
return x*f;
}
void add(int x,int y){
go[++tot]=y;
next[tot]=h[x];
h[x]=tot;
}
void travel(int x,int y){
a[++top]=x;
int t=h[x];
size[x]=1;
while (t){
if (!bz[go[t]]&&go[t]!=y) {
travel(go[t],x);
size[x]+=size[go[t]];
}
t=next[t];
}
}
void dfs(int x,int y){
int t=h[x];
while (t){
if (!bz[go[t]]&&go[t]!=y){
mx[go[t]]=max(mx[x],v[go[t]]);
d[go[t]]=(d[x]+v[go[t]])%p;
dfs(go[t],x);
}
t=next[t];
}
}
bool cmp(int x,int y){
return mx[x]<mx[y]||mx[x]==mx[y]&&x<y;
}
void calc(int f,int x){
sort(a+1,a+top+1,cmp);
int i,j,t;
fo(i,1,top){
t=a[i];
cnt[d[t]]++;
j=(d[t]+v[x])%p;
j=((j-mx[t])%p+p)%p;
j=(p-j)%p;
ans+=(ll)cnt[j]*f;
}
fo(i,1,top) cnt[d[a[i]]]--;
}
void solve(int x,int y){
top=0;
travel(x,0);
if (y) calc(-1,y);
int i,j=x,k=0,t;
while (1){
t=h[j];
while (t){
if (!bz[go[t]]&&go[t]!=k&&size[go[t]]>top/2){
k=j;
j=go[t];
break;
}
t=next[t];
}
if (!t) break;
}
mx[j]=v[j];
d[j]=0;
dfs(j,0);
calc(1,j);
bz[j]=1;
t=h[j];
while (t){
if (!bz[go[t]]) solve(go[t],j);
t=next[t];
}
}
int main(){
freopen("path.in","r",stdin);freopen("path.out","w",stdout);
n=read();p=read();
fo(i,1,n-1){
j=read();k=read();
add(j,k);add(k,j);
}
fo(i,1,n) v[i]=read();
solve(1,0);
printf("%lld\n",ans);
}