洛谷P3252 [JLOI2012]树

链接

　　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;
}