链接
https://www.luogu.org/problem/show?pid=3252
题解
如果权值都是正数的话,那很容易,但这道题的后6个点权值有负数出现(我被坑了)。
设sum[x]表示从根节点到节点x的权值和,那么一条路径的长度等于S,就等价于sum的差等于S,即sum[x]-sum[y]=S指的是点x和y的儿子构成长度为S的链。
变形一下,得到sum[x]=sum[y]+x。
从根节点开始dfs,那么当前正在访问的点恰好拉成一条链,如果我当前在点x,你要看当前链上有多少点y,其sum[y]+S恰好等于sum[x],计入答案就好。
和很显然就可以用权值树状数组了,离散然后做就行了。时间复杂度虽然是
O(NlogN)
,但在bzoj上依然484ms过。(这可能因为是树状数组常数小)
代码
//树上前缀和、树状数组
#include <cstdio>
#include <algorithm>
#define ll long long
#define maxn 300000
#define lowbit(x) (x&-x)
using namespace std;
ll N, S, w[maxn], ans, head[maxn], to[maxn], s[maxn], nex[maxn], tot, tmp[maxn],
ta[maxn];
void adde(ll a, ll b){to[++tot]=b;nex[tot]=head[a];head[a]=tot;}
void add(ll pos, ll v){for(;pos<=*tmp;pos+=lowbit(pos))ta[pos]+=v;}
ll read(ll x=0)
{
char c=getchar();
while(c<48 or c>57)c=getchar();
while(c>=48 and c<=57)x=(x<<3)+(x<<1)+c-48,c=getchar();
return x;
}
ll sum(ll pos)
{
ll ans=0;
for(;pos;pos-=lowbit(pos))ans+=ta[pos];
return ans;
}
void init()
{
ll i, x, y;
N=read(), S=read();
for(i=1;i<=N;i++)w[i]=read();
for(i=1;i<N;i++)x=read(), y=read(), adde(x,y);
}
void lisan()
{
ll i, x=0;
tmp[++*tmp]=S;
for(i=1;i<=N;i++)tmp[++*tmp]=s[i],tmp[++*tmp]=s[i]+S;
sort(tmp+1,tmp+*tmp+1);
for(i=1;i<=*tmp;i++)if(tmp[i]!=tmp[i-1])tmp[++x]=tmp[i];
*tmp=x;
}
ll find(ll x){return lower_bound(tmp+1,tmp+tmp[0]+1,x)-tmp;}
void dfs1(ll pos, ll sum)
{
ll p;
s[pos]=sum+w[pos];
for(p=head[pos];p;p=nex[p])dfs1(to[p],s[pos]);
}
void dfs2(ll pos)
{
ll p, t=find(s[pos]);
ans+=sum(t)-sum(t-1);
add(find(s[pos]+S),1);
for(p=head[pos];p;p=nex[p])dfs2(to[p]);
add(find(s[pos]+S),-1);
}
int main()
{
init();
dfs1(1,0);
lisan();
add(find(S),1);
dfs2(1);
printf("%lld",ans);
return 0;
}