[Luogu3676]小清新数据结构题
题面戳我
题意:给一棵树,树上有点权,每次操作为修改一个点的点权,或者是询问以某个点为根时,每棵子树(以每个点为根,就有n棵子树)点权和的平方和。
\(n\le2*10^5\),保证答案在long long范围内
sol
我们设\(s_i\)表示以\(p\)为整棵树的根时,以\(i\)为根的子树的点权和。设\(Sum\)表示所有点的点权和,即\(Sum=\sum_{i=1}^{n}val_i\)。
所以这道题给出\(p\),就是要你求\(\sum_{i=1}^{n}s_i^2\)。
我们先看\(\sum_{i=1}^{n}s_i\)怎么求。
考虑每个点的点权对\(\sum_{i=1}^{n}s_i\)的贡献,可以发现,每个点被计算了\(dep_i+1\)次,也就是说\(\sum_{i=1}^{n}s_i=\sum_{i=1}^{n}val_i(dep_i+1)=\sum_{i=1}^{n}val_idep_i+Sum\)。前面那一坨是不是有点熟悉?【ZJOI2015】幻想乡战略游戏。
下文中为了方便描述,令\(calc(p)\)表示以\(p\)为根时的\(\sum_{i=1}^{n}val_idep_i\)
接下来我们考虑一下这个东西
这个可以形象地理解为,在每一对点对\((i,j)\)的路径上每一条边(刚好是\(dis(i,j)\)条边)上都加上\(val_ival_j\),然后求整棵树上的边权之和。
现在我们考虑每一条边上的权值,它应该等于它两侧连接的两坨树的点权和的乘积。而连接的这两坨树中,不论取哪个\(p\)为根,都有有且仅有一坨树会是一棵子树。所以这个权值会等于\(s_i(Sum-s_i)\)。所以
这同时也证明了不论取哪个\(p\)作为根,\(\sum_{i=1}^{n}s_i(Sum-s_i)\)都不会变。
令\(W=\sum_{i=1}^{n}s_i(Sum-s_i)\),可以先\(O(n)\)地\(DP\)出\(W\)的初值,然后就只要考虑一个点权修改对\(W\)的影响。
因为\(W=\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)\),若节点\(i\)的点权的变化量为\(\Delta v\),那么\(\Delta W=\Delta v\sum_{j=1}^{n}val_jdis(i,j)\),相当于\(\Delta v*calc(i)\),所以说一样地计算即可。
所以最终询问的答案就是:
code
#include<cstdio>
#include<algorithm>
using namespace std;
#define ll long long
const int N = 200005;
int gi()
{
int x=0,w=1;char ch=getchar();
while ((ch<'0'||ch>'9')&&ch!='-') ch=getchar();
if (ch=='-') w=0,ch=getchar();
while (ch>='0'&&ch<='9') x=(x<<3)+(x<<1)+ch-'0',ch=getchar();
return w?x:-x;
}
struct edge{int to,next;}a[N<<1];
int n,q,val[N],head[N],cnt,pa[N],dep[N],sz[N],son[N],top[N];
void dfs1(int u,int f)
{
pa[u]=f;dep[u]=dep[f]+1;sz[u]=1;
for (int e=head[u];e;e=a[e].next)
{
int v=a[e].to;if (v==f) continue;
dfs1(v,u);
sz[u]+=sz[v];if (sz[v]>sz[son[u]]) son[u]=v;
}
}
void dfs2(int u,int f)
{
top[u]=f;
if (son[u]) dfs2(son[u],f);else return;
for (int e=head[u];e;e=a[e].next)
if (a[e].to!=pa[u]&&a[e].to!=son[u])
dfs2(a[e].to,a[e].to);
}
int lca(int u,int v)
{
while (top[u]^top[v])
{
if (dep[top[u]]<dep[top[v]]) swap(u,v);
u=pa[top[u]];
}
return dep[u]<dep[v]?u:v;
}
int getdis(int u,int v){return dep[u]+dep[v]-(dep[lca(u,v)]<<1);}
int tot,root,vis[N],w[N],fa[N];
ll sum[N],gather[N],tofa[N],sigma,omega,ans;
void getroot(int u,int f)
{
sz[u]=1;w[u]=0;
for (int e=head[u];e;e=a[e].next)
{
int v=a[e].to;if (v==f||vis[v]) continue;
getroot(v,u);
sz[u]+=sz[v];w[u]=max(w[u],sz[v]);
}
w[u]=max(w[u],tot-sz[u]);
if (w[u]<w[root]) root=u;
}
void solve(int u,int f)
{
fa[u]=f;vis[u]=1;
for (int e=head[u];e;e=a[e].next)
{
int v=a[e].to;if (vis[v]) continue;
tot=sz[v];
root=0;
getroot(v,0);
solve(root,u);
}
}
void modify(int u,int v)
{
sum[u]+=v;
for (int i=u;fa[i];i=fa[i])
{
int dist=getdis(u,fa[i]);
sum[fa[i]]+=v;
gather[fa[i]]+=dist*v;
tofa[i]+=dist*v;
}
}
ll calc(int u)
{
ll res=gather[u];
for (int i=u;fa[i];i=fa[i])
{
int dist=getdis(u,fa[i]);
res+=(ll)dist*(sum[fa[i]]-sum[i]);
res+=gather[fa[i]]-tofa[i];
}
return res;
}
void DP(int u)
{
sz[u]=val[u];
for (int e=head[u];e;e=a[e].next)
{
int v=a[e].to;if (v==pa[u]) continue;
DP(v);sz[u]+=sz[v];
}
omega+=1ll*sz[u]*(sigma-sz[u]);
}
int main()
{
n=gi();q=gi();
for (int i=1;i<n;i++)
{
int u=gi(),v=gi();
a[++cnt]=(edge){v,head[u]};head[u]=cnt;
a[++cnt]=(edge){u,head[v]};head[v]=cnt;
}
dfs1(1,0);dfs2(1,1);
tot=w[0]=n;
getroot(1,0);
solve(root,0);
for (int i=1;i<=n;i++)
val[i]=gi(),modify(i,val[i]),sigma+=val[i];
DP(1);
while (q--)
{
int opt=gi(),x=gi();
if (opt==1)
{
int y=gi();
modify(x,y-val[x]);sigma+=y-val[x];
omega+=(y-val[x])*calc(x);
val[x]=y;
}
else printf("%lld\n",(calc(x)+sigma)*sigma-omega);
}
return 0;
}