Luogu P3676 小清新数据结构题

Problem

一棵有点权树，每个节点都有一个s表示其子树权值和。支持点权修改，以及查询x为根时，所有节点s的平方和。

Solution

（一）考虑没有换根，即对于所有查询均有x=1
点修改u的权值，设增量为a。根据树的性质，只有path(1, u)上的点的s会被更改。
考虑这次修改对答案的贡献。
$\Rightarrow\sum_{i \in path(1,u)} (s_{i}+a)^{2}-\sum_{i \in path(1,u)} s_{i}^{2}$
$\Rightarrow a^{2} \cdot len + 2a \sum _{i \in path(1,u)}s_{i}$
于是每次点修改只需要更改一下路径（打个lazy tag），然后顺便更新ans。
询问时，ans就是答案。
（二）考虑换根（从1换到u）
设path(1, u)上的节点依次为$v_{1},v_{2}……v_{k}$
$v_{i}$在root=1下的子树点权为$a_{i}$，在root=u下的子树点权为$b_{i}$
$\star a_{i+1} + b_{i} = a_{1} = b_{k}$
开始推导。

$$\begin{aligned} ans_{u} &= ans_{1} - \sum_{i=1}^{k}a_{i}^{2}+\sum_{i=1}^{k}b_{i}^{2} \\ &=ans_{1}-\sum_{i=1}^{k}a_{i}^{2}+\sum_{i=1}^{k-1}(a_{1}-a_{i+1})^{2}+b_{k}^{2} \\ &=ans_{1} - \sum_{i=2}^{k}a_{i}^{2}+\sum_{i=2}^{k}(a_{1}-a_{i})^{2} \\ &=ans_{1} + (k - 1)a_{1}^{2}-2a_{1}\sum_{i=2}^{k}a_{i}\\ &=ans+s_{1} \cdot [(k-1)s_{1} - 2[(\sum_{i \in path(1, u)}s_{i}) - s_{1}]] \\ &=ans+s_{1} \cdot [(k+1)s_{1} - 2 \sum_{i \in path(1,u)}s_{i}] \end{aligned}$$
$ans$一直都是 $ans_{1}$。
s其实就是a。在做的时候其实并不换根，维护的全都是root=1，如果查询别的根的答案，就在目前ans基础上改动一下输出即可（并不改ans）。
可以计算出 $ans_{Max} < 3 \times 10^{17}$，用long long存储即可。

#include <bits/stdc++.h>
using namespace std;
#define N 200010
typedef long long ll;
inline char gc() {
    static char now[1<<16], *S, *T;
    if(S == T) {T = (S = now) + fread(now, 1, 1<<16, stdin); if(S == T) return EOF;}
    return *S++;
}
inline int read() {
    int x = 0, f = 1; char c = gc();
    while(c < '0' || c > '9') {if(c == '-') f = -1; c = gc();}
    while(c >= '0' && c <= '9') {x = x * 10 + c - 48; c = gc();}
    return x * f;
}
struct edge {int to, next;}e[N<<1];
ll head[N], c[N], w[N], fa[N], dep[N], sz[N], son[N], top[N], id[N], pos[N];
ll ans = 0;
ll n, m, cnt = 1, dfn = 0;
inline void ins(int x, int y) {e[++cnt].to = y; e[cnt].next = head[x]; head[x] = cnt;}
void dfs1(int x) {
    dep[x] = dep[fa[x]] + (sz[x] = 1); w[x] = c[x];
    for(int i = head[x]; i; i = e[i].next) {
        int y = e[i].to;
        if(y^fa[x]) {
            fa[y] = x; dfs1(y); w[x]+= w[y]; sz[x]+= sz[y];
            if(sz[y] > sz[son[x]]) son[x] = y;
        }
    }
    ans+= (ll)w[x] * w[x];
}
ll c1[N], c2[N];
inline void add(int x, ll w) {for(int i = x; i <= n; i+= i & -i) c1[i]+= w, c2[i]+= x * w;}
inline ll qry(int x) {ll w = 0; for(int i = x; i; i-= i & -i) w+= (ll)(x + 1) * c1[i] - c2[i]; return w;}
void dfs2(int x, int tp) {
    top[x] = tp; pos[id[x] = ++dfn] = x;
    add(dfn, w[x] - w[pos[dfn - 1]]);
    if(son[x]) dfs2(son[x], tp);
    for(int i = head[x]; i; i = e[i].next) if(e[i].to^fa[x] && e[i].to^son[x]) dfs2(e[i].to, e[i].to);
}
inline void modify(int u, int x) {
    int k = 0; ll s = 0;
    for(; u; u = fa[top[u]]) {
        k+= id[u] - id[top[u]] + 1;
        s+= qry(id[u]) - qry(id[top[u]] - 1), add(id[top[u]], x), add(id[u] + 1, -x);
    }
    ans+= (ll)x * (x * k + (s<<1));
}
inline ll query(int u) {
    int k = 0, x = qry(1); ll s = 0;
    for(; u; u = fa[top[u]]) {
        k+= id[u] - id[top[u]] + 1;
        s+= qry(id[u]) - qry(id[top[u]] - 1);
    }
    return ans + x * ((k + 1) * x - (s<<1));
}
int main() {
    n = read(); m = read();
    for(int i = 1; i < n; ++i) {int u = read(), v = read(); ins(u, v); ins(v, u);}
    for(int i = 1; i <= n; ++i) c[i] = read();
    dfs1(1); dfs2(1, 1);
    for(int i = 1; i <= m; ++i) {
        int o = read(), u = read(), v;
        if(o == 1) v = read(), v-= c[u], c[u]+= v, modify(u, v);
        else printf("%lld\n", query(u));
    }
    return 0;
}