luogu P3676 小清新数据结构题

首先我们考虑如果题目要求的不是子树平方和而是子树和，也就是\(\sum_{i=1}^{n}siz[i]\)，那么每个点的贡献就是他到根的距离加一，所以答案也可以写成：
\(\sum_{i=1}^{n}a[i] \times dis(i,p)+Sum\)（其中\(p\)表示当前的根，\(a[i]\)表示点\(i\)的点权，\(Sum=\sum a[i]\)）。
我们发现这个东西其实就是P3345 [ZJOI2015]幻想乡战略游戏这个题，为了后文柿子的表达方便，我们记
\(calc(p)=\sum_{i=1}^{n}a[i] \times dis(i,p)\)

然后这个平方很不好处理。首先给出一个结论，无论根在哪里，\(\sum_{i=1}^{n}siz[i] \times (Sum-siz[i])\)的值恒定（记这个常量为\(W\)）。

怎么证明呢，发现任意两点\(i,j\)对上述式子的贡献是\(a[i] \times a[j] \times dis(i,j)\)，这个东西无论根在哪里都不会变的。

然后稍微变形一下就能推出\(\sum siz[i]^2=Sum \times (calc(p)+Sum) - W\)。所以我们可以在根为\(1\)的时候求出\(W\)然后动态点分治维护\(calc(p)\)就好了。

现在的问题就是如何维护。除了\(W\)以外其他都是非常好维护的，我们发现一个节点\(k\)对\(W\)的贡献为
\(a[k] \times \sum_{i=1}^{n}a[i] \times dis(i,k)=a[k] \times calc(k)\)，然后对于怎么维护\(W\)就不用多说了吧。。。

代码：

#pragma GCC optimize(2)
#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#define rint int

using namespace std;

typedef long long LL;
const int N=1000009;
int n,Q,root,head[N],cnt,F[N],log[N],a[N];
LL dis1[N],dis2[N],sum[N],W,Sum;
struct Edge
{
	int nxt,to,w;
}g[N*2];
struct G
{
	int head[N],cnt,f[N][30],dep[N],del[N],siz[N],Euler[N],Fst[N],Index;
	struct Edge
	{
		int nxt,to;
	}g[N*2];
	
	void add(int from,int to)
	{
		g[++cnt].nxt=head[from];
		g[cnt].to=to;
		head[from]=cnt;
	}
	int LCA(int x,int y)
	{
		if(Fst[x]>Fst[y])
			swap(x,y);
		x=Fst[x],y=Fst[y];
		int k=log[y-x+1];
		return dep[f[x][k]]<dep[f[y-(1<<k)+1][k]]?f[x][k]:f[y-(1<<k)+1][k];
	}
	void dfs(int x,int fa)
	{
		Euler[++Index]=x,Fst[x]=Index,f[Index][0]=x;
		for (rint i=head[x];i;i=g[i].nxt)
		{
			int v=g[i].to;
			if(v==fa)
				continue;
			dep[v]=dep[x]+1;
			dfs(v,x);
			Euler[++Index]=x,f[Index][0]=x;
		}
	}
	int Get_dis(int x,int y)
	{
		return dep[x]+dep[y]-2*dep[LCA(x,y)];
	}
	void DFS(int x,int fa)
	{
		siz[x]=1;
		for (rint i=head[x];i;i=g[i].nxt)
		{
			int v=g[i].to;
			if(v==fa||del[v])
				continue;
			DFS(v,x);
			siz[x]+=siz[v];
		}
	}
	int Get_Weight(int x)
	{
		DFS(x,-1);
		int k=siz[x]/2,fa=-1;
		while(1)
		{
			int tmp=0;
			for (rint i=head[x];i;i=g[i].nxt)
			{
				int v=g[i].to;
				if(v==fa||del[v])
					continue;
				if(siz[v]>siz[tmp])
					tmp=v;
			}
			if(siz[tmp]<=k)
				return x;
			fa=x,x=tmp;
		}
	}
	void work()
	{
		dfs(1,-1);
		for (int j=1;1<<j<=Index;j++)
			for (int i=1;i+(1<<j)-1<=Index;i++)
				if(dep[f[i][j-1]]<dep[f[i+(1<<j-1)][j-1]])
					f[i][j]=f[i][j-1];
				else
					f[i][j]=f[i+(1<<j-1)][j-1];
	}
	void dfs_1(int x,int fa)
	{
		siz[x]=a[x];
		for (int i=head[x];i;i=g[i].nxt)
		{
			int v=g[i].to;
			if(v==fa)
				continue;
			dfs_1(v,x);
			siz[x]+=siz[v];
		}
	}
}T;

void add(int from,int to,int w)
{
	g[++cnt].nxt=head[from];
	g[cnt].to=to;
	g[cnt].w=w;
	head[from]=cnt;
}

int read()
{
	char c=getchar();
	int x=0,f=1;
	while(c<'0'||c>'9')
	{
		if(c=='-')
			f=-1;
		c=getchar();
	}
	while(c>='0'&&c<='9')
		x=x*10+c-'0',c=getchar();
	return x*f;
}

void init()
{
	log[0]=-1;
	for (int i=1;i<=N-9;i++)
		log[i]=log[i>>1]+1;
	n=read(),Q=read();
	for (rint i=1;i<n;i++)
	{
		int x,y;
		x=read(),y=read();
		T.add(x,y),T.add(y,x);
	}
	T.work();
}

void build(int fa)
{
	T.del[fa]=1;
	for (rint i=T.head[fa];i;i=T.g[i].nxt)
	{
		int v=T.g[i].to;
		if(T.del[v]||v==fa)
			continue;
		int w=T.Get_Weight(v);
		F[w]=fa,add(fa,w,v),add(w,fa,v);
		build(w);
	}
}

void Modify(int x,int y)
{
	for (rint i=x;i;i=F[i])
	{
		sum[i]+=y;
		if(i!=root)
		{
			LL fuck=1LL*T.Get_dis(x,F[i])*y;
			dis1[F[i]]+=fuck,dis2[i]+=fuck;
		}
	}
}

LL calc(int x)
{
	LL ans=dis1[x];
	for (rint i=x;F[i];i=F[i])
	{
		ans+=dis1[F[i]]-dis2[i];
		ans+=(sum[F[i]]-sum[i])*T.Get_dis(x,F[i]);
	}
	return ans;
}

void work()
{
	root=T.Get_Weight(1);
	build(root);
	for (int i=1;i<=n;i++)
		scanf("%d",&a[i]),Modify(i,a[i]),Sum+=a[i];
	T.dfs_1(1,-1);
	for (int i=1;i<=n;i++)
		W+=(Sum-T.siz[i])*T.siz[i];
	for (rint _=1;_<=Q;_++)
	{
		int x,y,opt;
		opt=read();
		if(opt==1)
		{
			x=read(),y=read();
			int delx=y-a[x];
			W+=calc(x)*delx,Sum+=delx;
			Modify(x,delx);
			a[x]=y;
		}
		else
		{
			x=read();
			printf("%lld\n",Sum*(Sum+calc(x))-W);
		}
	}
}

int main()
{
	init();
	work();
	return 0;
}