洛谷P3676 小清新数据结构题 【树剖 + BIT】
题目链接
题解
我们先维护\(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;
}