题目描述:
给一棵n个节点的带点权的树,求:
∑x=1n∑y=x+1n∑dis(x,y)+1k=1(路径x−y上前k大点的点权和)
∑
x
=
1
n
∑
y
=
x
+
1
n
∑
k
=
1
d
i
s
(
x
,
y
)
+
1
(
路
径
x
−
y
上
前
k
大
点
的
点
权
和
)
解题思路:
一个节点的权值会被计算所有链上点权小等于它的点数次。
我们把点按权值从小到大排序后依次插入,对于插入的点
u
u
,分情况讨论之前插入的一个点对
u
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;
}