【刷题】洛谷 P3676 小清新数据结构题

题目背景

本题时限2s，内存限制256M

题目描述

在很久很久以前，有一棵n个点的树，每个点有一个点权。

现在有q次操作，每次操作是修改一个点的点权或指定一个点，询问以这个点为根时每棵子树点权和的平方和。

（题目不是很好懂，没看太懂的可以看看样例解释）

输入输出格式

输入格式：

第一行两个整数n、q。

接下来n-1行每行两个整数a和b，表示树中a与b之间有一条边，保证给出的边不会重复。

接下来一行n个整数，第i个整数表示第i个点的点权。

接下来q行每行两或三个数，如果第一个数为1，那么接下来有两个数x和y，表示将第x个点的点权修改为y，如果第一个数为2，那么接下来有一个数x，表示询问以x为根时每棵子树点权和的平方和。

输出格式：

对于每个询问输出答案。

输入输出样例

输入样例#1：

4 5

1 2

2 3

2 4

4 3 2 1

2 2

1 1 3

2 3

1 2 4

2 4

输出样例#1：

121

140

194

说明

样例解释

这是一个菊花图，2与1、3、4间有边。

一开始每个点点权分别为4、3、2、1。

第一个询问以2为根，1、3、4子树中都只有本身一个点，2子树中有所有点，那么1、3、4子树中的点权和就分别是自己的点权4、2、1，2子树中的点权和就是4+3+2+1=10， \(4^2+2^2+1^1+10^2=121\) 。

接下来将第一个点点权修改为3，每个点点权分别为3、3、2、1。

第二个询问以3为根，1、4子树中只有自己，2子树中有1、2、4，3子树中有所有点，1、4子树点权和就是自己的点权3、1，2子树点权和就是3+3+1=7，3子树点权和为3+3+2+1=9， \(3^2+1^2+7^2+9^2=140\) 。

接下来把第二个点点权修改为4，每个点点权分别为3、4、2、1。

第三个询问以4为根，1、3子树点权和就是3和2，2子树点权和就是3+4+2=9，4子树点权和为3+4+2+1=10， \(3^2+2^2+9^2+10^2=194\) 。

数据范围

对于10%的数据， \(n,q \leq 2000\) 。

对于40%的数据， \(n,q \leq 60000\) 。

对于另外30%的数据，保证每次询问的根都为1。

对于100%的数据， \(1 \leq n,q \leq 200000\) ， \(-10 \leq\) 输入的每个点权 \(\leq 10\) 。

建议使用输入优化，虽然标程没加读入优化也能过

题解

不会动态点分治，于是树剖乱搞

树剖加线段树，线段树维护一个区间和，一个区间和的平方，一个lazy标记

考虑部分分，所有询问的根均为1的怎么做

如果一个修改改的是 \(u\) 点，那么这次修改影响了的点只有 \(u\) 到 \(1\) 号点路径上的点，只有这些点的子树和和子树和的平方有变化，考虑这个变化

令 \(p=newval_u-oldval_u\)，即 \(p\) 是 \(u\) 点增加（或减少）的权值

它到根的路径上的点的子树和修改很简单，直接线段树区间add就行了，看和的平方的增加，把剖好后的每一条链对应到区间上， \(\sum (s_i+p)^2=\sum s_i^2+2s_ip+p^2=\sum s_i^2+2p\sum s_i+p^2len\) ，可以知道这一段区间所有位置的平方和的增加的贡献是一样的，因此也可以lazy。

那么如果一个区间上有lazy标记 \(p\) ，先要算新的区间和的平方（因为式子与旧的区间和有关系，如果先算区间和会导致平方的答案错误），直接加 \(2p\sum s_i+p^2len\) ，然后算新的区间和

然后询问的根为1的就做完了，答案就是整个线段树的区间和

然后看询问的根不为1的怎么做

同样，默认把根作为1，按照根为1的方法维护线段树，可以得到某棵树中，根为1的答案 \(stans\)

然后就看，怎样把根为1的答案变成根为询问的点的答案

假设询问的点为 \(u\)

同样，树的形态改变后，子树和有变化的也只有 \(u\) 到 \(1\) 的路径上的点，将这些点按深度从小到大排序，得到序列 \(p_1,p_2,p_3,...,p_{cnt}\)

易得式子， \(ans=stans-\sum_{i=1}^{cnt} a_{p_i}^2+\sum_{i=1}^{cnt}b_{p_i}^2\)，其中 \(a_i\) 表示以1为根时， \(i\) 的子树和， \(b_i\) 表示以 \(u\) 为根时， \(i\) 的子树和

然后有这样一个等式，\(b_{p_i}+a_{p_{i+1}}=a_{p_1}=b_{p_{cnt}}\) ，它们都等于整棵树的权值和的平方

然后式子化简

\(ans=stans-\sum_{i=1}^{cnt} a_{p_i}^2+b_{p_{cnt}}^2+\sum_{i=1}^{cnt-1}(a_{p_1}-a_{p_{i+1}})^2\)

\(~~~~~~~~=stans+a_{p_1}^2-\sum_{i=1}^{cnt} a_{p_i}^2+\sum_{i=1}^{cnt-1}a_{p_1}^2-2a_{p_1}a_{p_{i+1}}+a_{p_{i+1}}^2\)

\(~~~~~~~~=stans+a_{p_1}^2-\sum_{i=1}^{cnt} a_{p_i}^2+(cnt-1)a_{p_1}^2-2a_{p_1}\sum_{i=2}^{cnt}a_{p_i}+\sum_{i=2}^{cnt}a_{p_i}^2\)

\(~~~~~~~~=stans+a_{p_1}^2-a_{p_1}^2+(cnt-1)a_{p_1}^2-2a_{p_1}\sum_{i=1}^{cnt}a_{p_i}+2a_{p_1}^2\)

\(~~~~~~~~=stans+(cnt+1)a_{p_1}^2-2a_{p_1}\sum_{i=1}^{cnt}a_{p_i}\)

而 \(cnt\) 其实就是 \(u\) 的dep，\(\sum_{i=1}^{cnt}a_{p_i}\) 就是以1为根的时候 \(u\) 到 \(1\) 路径上的子树和的和，这个就在线段树里已经维护了

然后直接算新的答案就可以了，程序里维护的树的形态根本不用改变

#include<bits/stdc++.h>

#define ui unsigned int

#define ll long long

#define db double

#define ld long double

#define ull unsigned long long

const int MAXN=200000+10,inf=0x3f3f3f3f;

int n,q,e,to[MAXN<<1],nex[MAXN<<1],beg[MAXN],sum[MAXN],vl[MAXN],st[MAXN],ed[MAXN],cnt,fa[MAXN],top[MAXN],hson[MAXN],size[MAXN],val[MAXN],dep[MAXN];

ll stans;

#define Mid ((l+r)>>1)

#define ls rt<<1

#define rs rt<<1|1

#define lson ls,l,Mid

#define rson rs,Mid+1,r

struct Segment_Tree{

	ll Sum1[MAXN<<2],Sum2[MAXN<<2],Add[MAXN<<2];

	inline void PushUp(int rt)

	{

		Sum1[rt]=Sum1[ls]+Sum1[rs];

		Sum2[rt]=Sum2[ls]+Sum2[rs];

	}

	inline void PushDown(int rt,int len)

	{

		Sum2[ls]+=1ll*Add[rt]*Add[rt]*(len-(len>>1))+2ll*Add[rt]*Sum1[ls];

		Sum2[rs]+=1ll*Add[rt]*Add[rt]*(len>>1)+2ll*Add[rt]*Sum1[rs];

		Sum1[ls]+=1ll*Add[rt]*(len-(len>>1));

		Sum1[rs]+=1ll*Add[rt]*(len>>1);

		if(Add[ls]==inf)Add[ls]=0;

		if(Add[rs]==inf)Add[rs]=0;

		Add[ls]+=Add[rt];Add[rs]+=Add[rt];

		Add[rt]=inf;

	}

	inline void Build(int rt,int l,int r)

	{

		Add[rt]=inf;

		if(l==r)

		{

			Sum1[rt]=vl[l];

			Sum2[rt]=1ll*vl[l]*vl[l];

		}

		else

		{

			Build(lson);Build(rson);

			PushUp(rt);

		}

	}

	inline void Update(int rt,int l,int r,int L,int R,ll k)

	{

		if(L<=l&&r<=R)

		{

			Sum2[rt]+=1ll*k*k*(r-l+1)+2ll*k*Sum1[rt];

			Sum1[rt]+=1ll*(r-l+1)*k;

			if(Add[rt]==inf)Add[rt]=0;

			Add[rt]+=k;

		}

		else

		{

			if(Add[rt]!=inf)PushDown(rt,r-l+1);

			if(L<=Mid)Update(lson,L,R,k);

			if(R>Mid)Update(rson,L,R,k);

			PushUp(rt);

		}

	}

	inline ll Query(int rt,int l,int r,int L,int R)

	{

		if(L<=l&&r<=R)return Sum1[rt];

		else

		{

			ll res=0;

			if(Add[rt]!=inf)PushDown(rt,r-l+1);

			if(L<=Mid)res+=Query(lson,L,R);

			if(R>Mid)res+=Query(rson,L,R);

			return res;

		}

	}

};

Segment_Tree T;

#undef Mid

#undef ls

#undef rs

#undef lson

#undef rson

template<typename T> inline void read(T &x)

{

	T data=0,w=1;

	char ch=0;

	while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();

	if(ch=='-')w=-1,ch=getchar();

	while(ch>='0'&&ch<='9')data=((T)data<<3)+((T)data<<1)+(ch^'0'),ch=getchar();

	x=data*w;

}

template<typename T> inline void write(T x,char ch='\0')

{

	if(x<0)putchar('-'),x=-x;

	if(x>9)write(x/10);

	putchar(x%10+'0');

	if(ch!='\0')putchar(ch);

}

template<typename T> inline void chkmin(T &x,T y){x=(y<x?y:x);}

template<typename T> inline void chkmax(T &x,T y){x=(y>x?y:x);}

template<typename T> inline T min(T x,T y){return x<y?x:y;}

template<typename T> inline T max(T x,T y){return x>y?x:y;}

inline void insert(int x,int y)

{

	to[++e]=y;

	nex[e]=beg[x];

	beg[x]=e;

}

inline void dfs(int x,int f)

{

	sum[x]=val[x];

	for(register int i=beg[x];i;i=nex[i])

		if(to[i]==f)continue;

		else dfs(to[i],x),sum[x]+=sum[to[i]];

}

inline void dfs1(int x,int f)

{

	int res=0;

	size[x]=1;fa[x]=f;dep[x]=dep[f]+1;

	for(register int i=beg[x];i;i=nex[i])

		if(to[i]==f)continue;

		else

		{

			dfs1(to[i],x);

			size[x]+=size[to[i]];

			if(size[to[i]]>res)res=size[to[i]],hson[x]=to[i];

		}

}

inline void dfs2(int x,int tp)

{

	st[x]=++cnt;vl[cnt]=sum[x];top[x]=tp;

	if(hson[x])dfs2(hson[x],tp);

	for(register int i=beg[x];i;i=nex[i])

		if(to[i]==fa[x]||to[i]==hson[x])continue;

		else dfs2(to[i],to[i]);

	ed[x]=cnt;

}

inline void init()

{

	dfs1(1,0);dfs2(1,1);

	T.Build(1,1,n);stans=T.Sum2[1];

}

inline void Doans(int x,ll k)

{

	while(x)

	{

		T.Update(1,1,n,st[top[x]],st[x],k);

		x=fa[top[x]];

	}

}

inline ll Getans(int x)

{

	if(x==1)return stans;

	ll a0=T.Query(1,1,n,st[1],st[1]),res=stans+1ll*(dep[x]+1)*a0*a0,nsum=0;

	while(x)

	{

		nsum+=T.Query(1,1,n,st[top[x]],st[x]);

		x=fa[top[x]];

	}

	return res-2ll*a0*nsum;

}

int main()

{

	read(n);read(q);

	for(register int i=1;i<n;++i)

	{

		int u,v;read(u);read(v);

		insert(u,v);insert(v,u);

	}

	for(register int i=1;i<=n;++i)read(val[i]);

	dfs(1,0);

	init();

	while(q--)

	{

		int opt;read(opt);

		if(opt==1)

		{

			int x,y;read(x);read(y);

			Doans(x,y-val[x]);

			val[x]=y;stans=T.Sum2[1];

		}

		if(opt==2)

		{

			int x;read(x);

			write(Getans(x),'\n');

		}

	}

	return 0;

}

巴特西