洛谷P3676 小清新数据结构题 【树剖 + BIT】

题目链接

洛谷P3676

题解

我们先维护\(1\)为根的答案，再考虑换根

一开始的答案可以\(O(n)\)计算出来

考虑修改，记\(s[u]\)表示\(u\)为根的子树的权值和

当\(u\)节点产生\(v\)的增量时，只影响\(1\)到\(u\)路径上的\(s\)，权值和都\(+v\)

而对答案的影响是

\[\begin{aligned}
\Delta ans &= \sum\limits_{i}^{k}(s_i + v)^{2} - \sum\limits_{i = 1}^{k} s_i^2 \\
&= 2v\sum\limits_{i = 1}^{k}s_i + v^{2}k
\end{aligned}
\]

所以只需维护路径上\(s\)的和即可更新答案

使用树剖 + BIT或者线段树

考虑如何换根

考虑从\(1\)到\(u\)经历了什么

发现同样只有\(1\)到\(u\)路径上的点的贡献产生了变化

设\(a_i\)表示当\(1\)为根时，路径上第\(i\)个点的子树权值和

设\(b_i\)表示当\(u\)为根时，路径上第\(i\)个点的子树权值和

有等式

\[a_{i + 1} + b_i = a_1 = b_k
\]

而

\[\begin{aligned}
ans_u &= ans_1 - \sum\limits_{i = 1}^{k}a_i^{2} + \sum\limits_{i = 1}^{k}b_i^{2} \\
&= ans_1 - \sum\limits_{i = 1}^{k}a_i^{2} + \sum\limits_{i = 2}^{k}(a_1 - a_i)^{2} + a_1^2 \\
&= ans_1 - \sum\limits_{i = 2}^{k}a_i^{2} + \sum\limits_{i = 2}^{k}(a_1 - a_i)^{2} \\
&= ans_1 + (k - 1)a_1^{2} - 2a_1\sum\limits_{i = 2}^{k}a_i \\
&= ans_1 + s_1((k - 1)s_1 - 2(\sum\limits_{i = 1}^{k}s_i - s_1)) \\
&= ans_1 + s_1((k + 1)s_1 - 2\sum\limits_{i = 1}^{k}s_i)
\end{aligned}
\]

同样对\(\sum\limits_{i = 1}^{k}s_i\)求和就行了

复杂度\(O(nlog^2n)\)

\(BIT\)大法好

#include<algorithm>
#include<iostream>
#include<cstdio>
#define LL long long int
#define REP(i,n) for (int i = 1; i <= (n); i++)
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define lbt(x) (x & -x)
using namespace std;
const int maxn = 200005;
inline int read(){
	int out = 0,flag = 1; char c = getchar();
	while (c < 48 || c > 57){if (c == '-') flag = 0; c = getchar();}
	while (c >= 48 && c <= 57){out = (out << 1) + (out << 3) + c - 48; c = getchar();}
	return flag ? out : -out;
}
int n,Q,h[maxn],ne = 1;
int fa[maxn],top[maxn],siz[maxn],dfn[maxn],son[maxn],dep[maxn],cnt;
LL S[maxn],SS[maxn],s[maxn],val[maxn],ans;
struct EDGE{int to,nxt;}ed[maxn << 1];
inline void build(int u,int v){
	ed[++ne] = (EDGE){v,h[u]}; h[u] = ne;
	ed[++ne] = (EDGE){u,h[v]}; h[v] = ne;
}
void add(LL* s,int u,LL x){while (u <= n) s[u] += x,u += lbt(u);}
LL query(LL* s,int u){LL re = 0; while (u) re += s[u],u -= lbt(u); return re;}
void Add(int l,int r,LL v){
	add(S,l,v); add(S,r + 1,-v);
	add(SS,l,v * l); add(SS,r + 1,-v * (r + 1));
}
LL Query(int u){return 1ll * (u + 1) * query(S,u) - query(SS,u);}
LL Sum(int l,int r){return Query(r)  - Query(l - 1);}
void dfs1(int u){
	siz[u] = 1; s[u] = val[u];
	Redge(u) if ((to = ed[k].to) != fa[u]){
		fa[to] = u; dep[to] = dep[u] + 1; dfs1(to);
		siz[u] += siz[to]; s[u] += s[to];
		if (!son[u] || siz[son[u]] < siz[to]) son[u] = to;
	}
}
void dfs2(int u,int flag){
	dfn[u] = ++cnt; Add(cnt,cnt,s[u]);
	top[u] = flag ? top[fa[u]] : u;
	if (son[u]) dfs2(son[u],true);
	Redge(u) if ((to = ed[k].to) != fa[u] && to != son[u])
		dfs2(to,false);
}
void solve1(int u,LL v){
	LL len = dep[u] + 1,tot = 0,x = u; v -= val[u]; val[u] += v;
	while (x){
		tot += Sum(dfn[top[x]],dfn[x]);
		Add(dfn[top[x]],dfn[x],v);
		x = fa[top[x]];
	}
	ans += 2ll * v * tot + v * v * len;
}
void solve2(int u){
	LL tot = 0,s1 = Query(1),x = u,len = dep[u] + 1;
	while (x) tot += Sum(dfn[top[x]],dfn[x]),x = fa[top[x]];
	printf("%lld\n",ans + s1 * ((len + 1) * s1 - 2 * tot));
}
int main(){
	n = read(); Q = read();
	for (int i = 1; i < n; i++) build(read(),read());
	REP(i,n) val[i] = read();
	dfs1(1); dfs2(1,0);
	REP(i,n) ans += s[i] * s[i];
	int opt,x;
	while (Q--){
		opt = read(); x = read();
		if (opt & 1) solve1(x,read());
		else solve2(x);
	}
	return 0;
}