BJ模拟 链上求和【树状数组维护dfs序】

题目描述：

给一棵n个节点的带点权的树，求：
$\sum\limits_{x=1}^n\sum\limits_{y=x+1}^n\sum_{k=1}^{dis(x,y)+1}(路径x-y上前k大点的点权和)$

解题思路：

一个节点的权值会被计算所有链上点权小等于它的点数次。
我们把点按权值从小到大排序后依次插入，对于插入的点$u$，分情况讨论之前插入的一个点$v$对$u$的次数贡献($v在子树内或v是u的祖先$)，用树状数组维护即可，具体可以自己画图推推再结合代码。

#include<bits/stdc++.h>
#define ll long long
using namespace std;
int getint()
{
    int i=0,f=1;char c;
    for(c=getchar();(c!='-')&&(c<'0'||c>'9');c=getchar());
    if(c=='-')c=getchar(),f=-1;
    for(;c>='0'&&c<='9';c=getchar())i=(i<<3)+(i<<1)+c-'0';
    return i*f;
}
const int N=100005,mod=1e9+7;
int n,ans;
struct node{int val,id;}a[N];
inline bool operator < (const node &a,const node &b){return a.val<b.val;}
int idx,in[N],out[N],fa[N],size[N];
int tot,first[N],nxt[N<<1],to[N<<1];
struct Bit
{
    ll a[N];
    inline void add(int i,int v)
    {
        for(;i<=n;i+=i&-i)a[i]+=+v;
    }
    inline ll query(int i)
    {
        ll res=0;
        for(;i;i-=i&-i)res+=a[i];
        return res;
    }
    inline void add(int l,int r,int v)
    {
        add(l,v),add(r+1,-v);
    }
    inline ll query(int l,int r)
    {
        return query(r)-query(l-1);
    }
}bit1,bit2,bit3;
void add(int x,int y)
{
    nxt[++tot]=first[x],first[x]=tot,to[tot]=y;
}
void dfs(int u)
{
    in[u]=++idx,size[u]=1;
    for(int e=first[u];e;e=nxt[e])
    {
        int v=to[e];
        if(v==fa[u])continue;
        fa[v]=u;
        dfs(v),size[u]+=size[v];
    }
    out[u]=idx;
}
int calc(int u,int w)
{
    ll res=0,cnt=1;
    for(int e=first[u];e;e=nxt[e])
    {
        int v=to[e];
        if(v==fa[u])
        {
            ll tmp=bit1.query(n)-bit1.query(in[u],out[u]);
            tmp=tmp+bit3.query(in[u])*n-bit2.query(in[u]);
            res=(res+tmp*size[u])%mod;
            res=(res+cnt*(n-size[u]))%mod,cnt+=n-size[u];
        }
        else
        {
            res=(res+bit1.query(in[v],out[v])*(n-size[v]))%mod;
            res=(res+cnt*size[v])%mod,cnt+=size[v];
        }
    }
    bit1.add(in[u],size[u]);
    bit2.add(in[u],out[u],size[u]),bit3.add(in[u],out[u],1);
    for(int e=first[u];e;e=nxt[e])
    {
        int v=to[e];
        if(v!=fa[u])bit2.add(in[v],out[v],size[v]);
    }
    return res*w%mod;
}
int main()
{
    //freopen("lx.in","r",stdin);
    n=getint();
    int x,y;
    for(int i=1;i<n;i++)
    {
        x=getint(),y=getint();
        add(x,y),add(y,x);
    }
    for(int i=1;i<=n;i++)a[i]=(node){getint(),i};
    dfs(1),sort(a+1,a+n+1);
    for(int i=1;i<=n;i++)
        ans=(ans+calc(a[i].id,a[i].val))%mod;
    cout<<ans<<'\n';
    return 0;
}