[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\)

接下来我们考虑一下这个东西

\[\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)
\]

这个可以形象地理解为，在每一对点对\((i,j)\)的路径上每一条边（刚好是\(dis(i,j)\)条边）上都加上\(val_ival_j\)，然后求整棵树上的边权之和。

现在我们考虑每一条边上的权值，它应该等于它两侧连接的两坨树的点权和的乘积。而连接的这两坨树中，不论取哪个\(p\)为根，都有有且仅有一坨树会是一棵子树。所以这个权值会等于\(s_i(Sum-s_i)\)。所以

\[\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)=\sum_{i=1}^{n}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)\)，所以说一样地计算即可。

所以最终询问的答案就是：

\[\sum_{i=1}^{n}s_i^2=Sum*\sum_{i=1}^{n}s_i-W=Sum(calc(i)+Sum)-W
\]

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