XSY 2139 - 链上求和

题目精选 DFS序 树状数组 分析

---

题意

分析

首先，我们定义一个点的大小$w=(权值，编号)$。

$$\begin{aligned}&~~~~\sum_{x=1}^{n-1}\sum_{y=x+1}^n\sum_{k=1}^{dist(x,y)+1}(路径x-y上前k大点权)\\&=\sum_{x=1}^{n-1}\sum_{y=x+1}^n\sum_{u}w_u(\sum_{v}[w_u>w_v]+1)\\&=\sum_{1\le x<y\le n}\sum_uw_u\sum_{v}[w_u>w_v]+\sum_{1\leq x<y\le n}\sum_{u}w_u\end{aligned}$$
按照$w$从小到大排序，逐个插入，那么只用考虑之前插入的点对当前点的贡献，然后再考虑当前点对之后的点的贡献即可。

第二部分很方便维护，考虑解决第一部分。

对于每一个点$y$，考虑$y$对$x$的贡献。
①当$y$在$root-x$上，设$root--y-nxt--x$，则贡献为$val[x]*(n-siz[nxt])*siz[x]$
总和为$val[x]*siz[x]*\sum_y(n-siz[y])]$

②当$y$在$x$子树内，设$x-nxt-y$，则贡献为$val[x]*(n-siz[nxt])*siz[y]$
总和为$val[x]*(n-siz[nxt])*\sum_{y\in subtree(x)}siz[y]$
我们枚举$nxt$，维护子树内$siz$的和。

③当$y$在$x$的子树外，且不在$root-x$上时，贡献为$val[x]*siz[x]*siz[y]$
总和为$val[x]*siz[x]*(\sum_{y在x子树外} siz[y]-\sum_{y在root-x上}siz[x])$
用之前维护的东西求解即可。

代码

#include <cstdio>
#include <iostream>
#include <cctype>
#include <algorithm>
#include <vector>
using namespace std;
#define F(i,a,b) for (int i=(a);i<=(b);i++)
#define fore(v,x) for (vector<int>::iterator v=g[x].begin();v!=g[x].end();v++)
#define LL long long
const int N=131072;
const int M=262144;
const int L=(int)1e9+7;
int n,w[N];
vector<int> g[N];
struct Order {
    int w,id;
    Order(int _w=0,int _id=0):w(_w),id(_id){}
    friend int operator < (Order a,Order b) {
        if (a.w!=b.w) return a.w<b.w; else return a.id<b.id;
    }
}ord[N];
int s[N],e[N],ind;
int pre[N],siz[N];
int ans;
struct Tree {
    int c[M];
    int lowbit(int i) {
        return i&-i;
    }
    void Add(int x,int ad) {
        for (int i=x;i<=ind;i+=lowbit(i))
            (c[i]+=ad)%=L;
    }
    int Query(int x) {
        int sum=0;
        for (int i=x;i>0;i-=lowbit(i))
            (sum+=c[i])%=L;
        return sum;
    }
}c1,c2,c3;
namespace input {
    const int SIZE=1<<16|1;
    char buffer[SIZE];
    char *head=buffer+SIZE;
    const char *tail=head;
    char getch(void) {
        if (head==tail) {
            fread(buffer,1,SIZE,stdin);
            head=buffer;
        }
        return *head++;
    }
    int rd(void) {
        int x=0,f=1; char c=getch();
        for (;!isdigit(c);c=getch()) if (c=='-') f=-1;
        for (;isdigit(c);c=getch()) x=x*10+c-'0';
        return x*f;
    }
}
using input::rd;
void DFS(int x) {
    s[x]=++ind;
    siz[x]=1;
    fore(v,x)
        if (*v!=pre[x]) {
            pre[*v]=x;
            DFS(*v);
            siz[x]+=siz[*v];
        }
    e[x]=++ind;
}
int main(void) {
    #ifndef ONLINE_JUDGE
        freopen("sdchr.in","r",stdin);
        freopen("sdchr.out","w",stdout);
    #endif
    cin>>n;
    F(i,1,n-1) {
        int u=rd(),v=rd();
        g[u].push_back(v),g[v].push_back(u);
    }
    F(i,1,n) w[i]=rd();
    F(i,1,n) ord[i]=Order(w[i],i);
    sort(ord+1,ord+n+1);
    DFS(1);
    F(i,1,n) {
        //考虑点i 
        int s=1;
        fore(v,i)
            if (*v!=pre[i]) {
                (ans+=(LL)w[i]*s%L*siz[*v]%L)%=L;
                s+=siz[*v];
            }
        if (i!=1)
            (ans+=(LL)w[i]*s%L*(n-siz[i])%L)%=L;
    }
    F(i,1,n) {
        int x=ord[i].id;
        //到根链
        (ans+=(LL)w[x]*siz[x]%L*c1.Query(s[x])%L)%=L;
        //子树
        fore(v,x) if (*v!=pre[x]) {
            int t=(c2.Query(e[*v])-c2.Query(s[*v]-1))%L;
            (ans+=(LL)w[x]*(n-siz[*v])%L*t%L)%=L;
        }
        //子树外且非到根链 
        int t=(((c2.Query(ind)-(c2.Query(e[x])-c2.Query(s[x]-1)))%L)%L-c3.Query(s[pre[x]]))%L;
        (ans+=(LL)w[x]*siz[x]%L*t%L)%=L;
        fore(v,x) if (*v!=pre[x]) {
            c1.Add(s[*v],n-siz[*v]);
            c1.Add(e[*v],-(n-siz[*v]));
        }
        c2.Add(s[x],siz[x]);
        c3.Add(s[x],siz[x]);
        c3.Add(e[x],-siz[x]);
    }
    (ans+=L)%=L;
    cout<<ans<<endl;
    return 0;
}