【ACWing】2568. 树链剖分

题目地址：

https://www.acwing.com/problem/content/2570/

给定一棵树，树中包含$n$个节点（编号$1\sim n$），其中第$i$个节点的权值为$a_i$。初始时，$1$号节点为树的根节点。现在要对该树进行$m$次操作，操作分为以下$4$种类型：
1 u v k，修改路径上节点权值，将节点$u$和节点$v$之间路径上的所有节点（包括这两个节点）的权值增加$k$。
2 u k，修改子树上节点权值，将以节点$u$为根的子树上的所有节点的权值增加$k$。
3 u v，询问路径，询问节点$u$和节点$v$之间路径上的所有节点（包括这两个节点）的权值和。
4 u，询问子树，询问以节点$u$为根的子树上的所有节点的权值和。

输入格式：
第一行包含一个整数$n$，表示节点个数。
第二行包含$n$个整数，其中第 i 个整数表示$a_i$。
接下来$n-1$行，每行包含两个整数$x,y$，表示节点$x$和节点$y$之间存在一条边。
再一行包含一个整数$m$，表示操作次数。
接下来$m$行，每行包含一个操作，格式如题目所述。

输出格式：
对于每个操作$3$和操作$4$，输出一行一个整数表示答案。

数据范围：
$1\le n,m\le 10^5$
$0\le a_i,k\le 10^5$
$1\le u,v,x,y\le n$

先介绍几个概念：
1、重儿子，轻儿子，重链：重儿子指的是，对于每个节点，其所有子树中，节点最多的那棵子树的树根。如果有若干棵子树节点数一样，则任选一个；不是重儿子的节点称为轻儿子。重链指的是极大的由重边构成的路径。叶子自成一条重链。
2、重边，轻边：重儿子向上连的边称为重边，非重边的边称为轻边。

接着从树根开始DFS，每次优先走重儿子。遍历完成之后的先序序列里，可以保证重链上的点都是连续的一段。

有一个性质，是树链剖分的精髓所在。任意一个节点$u$到树根路径上最多有$O(\log n)$条重链或者轻边。如果$u$向上走的是一条轻边，那么说明$u$是轻儿子，轻儿子为根的子树的节点个数不超过$u$父亲子树的节点个数的一半，所以$u$每次沿着轻边向上走一步，$u$子树的节点个数会至少翻倍，从而其最多只能向上走$O(\log n)$条轻边。由于重链或者在轻边之间，或者在开始或结尾，所以重链数量也是$O(\log n)$级别。所以由这个性质，任意一条树中的路径，都是若干轻边和重链的合并。

本题的算法如下：
1、先从树根开始做一次DFS，此次DFS记录每个节点的深度、父亲，并且求一下每个节点的子树大小，和其重儿子。
2、再从树根开始做第二次DFS，此次DFS要将树做重链分解（即对每个顶点按DFS先序重新编号），记录新编号下每个点的权值，以及每个点所在重链的头结点（即深度最小的节点）。DFS的时候要优先遍历重儿子。
3、重链分解做好之后得到的DFS序列中，重链就都成为了区间。从而对任意路径的操作，都转化为了对区间的操作，从而可以用线段树来维护。具体说来就是：
（1）将$u$到$v$的路径上所有节点权值增加$k$：主要的思想是要找到它们的最近公共祖先。如果$u$和$v$不位于同一条重链上，那么比较一下它们所在的重链头结点，不妨设$u$的重链头结点更深，那么说明最近公共祖先在$u$重链头结点上方，从而可以直接将$u$一直到其重链头结点这段区间整体加$k$（在DFS序里重链就是一段区间，从而可以用线段树$O(\log n)$时间内进行区间操作），然后将$u$跳到重链头结点父亲处。接下来进行类似的操作，直到$u$和$v$爬到相同的重链上为止。这时候最近公共祖先就是$u$和$v$之一，再做最后一次区间操作即可。
（2）将$u$子树每个节点权值增加$k$。DFS序已经保证了子树一定是个区间了，并且树根一定是区间左端点，直接操作即可。
（3）询问$u$到$v$的路径和：同（1）
（4）询问$u$子树的权值和：同（2）

代码如下：

#include <iostream>
#include <cstring>
using namespace std;

const int N = 1e5 + 10, M = N * 2;
int n, m;
// w为节点权值
int h[N], w[N], e[M], ne[M], idx;
// id[x]为节点x的新编号，nw[x]是新编号为x的节点的权值
int id[N], nw[N], cnt;
// dep为深度，sz为子树大小，top[x]是x所在重链的头结点，
// fa[x]为x父亲，son[x]为x的重儿子
int dep[N], sz[N], top[N], fa[N], son[N];
struct Tree {
    int l, r;
    long sum, add;
} tr[N << 2];

void add(int a, int b) {
  e[idx] = b, ne[idx] = h[a], h[a] = idx++;
}

// 第一次dfs，求节点深度、父亲、子树大小和重儿子
void dfs1(int u, int from, int depth) {
  dep[u] = depth, fa[u] = from, sz[u] = 1;
  for (int i = h[u]; ~i; i = ne[i]) {
    int v = e[i];
    if (v == from) continue;
    dfs1(v, u, depth + 1);
    sz[u] += sz[v];
    if (sz[son[u]] < sz[v]) son[u] = v;
  }
}

// 第二次dfs，t为u重链头结点
void dfs2(int u, int t) {
  id[u] = ++cnt, nw[cnt] = w[u], top[u] = t;
  // 到叶子了，直接返回
  if (!son[u]) return;
  // 先遍历重儿子
  dfs2(son[u], t);
  // 遍历轻儿子
  for (int i = h[u]; ~i; i = ne[i]) {
    int v = e[i];
    if (v == fa[u] || v == son[u]) continue;
    dfs2(v, v);
  }
}

void pushup(int u) {
  tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
}

void pushdown(int u) {
  auto &root = tr[u], &left = tr[u << 1], &right = tr[u << 1 | 1];
  if (root.add) {
    left.sum += root.add * (left.r - left.l + 1);
    left.add += root.add;
    right.sum += root.add * (right.r - right.l + 1);
    right.add += root.add;
    root.add = 0;
  }
}

void build(int u, int l, int r) {
  tr[u] = {l, r, nw[l], 0};
  if (l == r) return;
  int mid = l + r >> 1;
  build(u << 1, l, mid), build(u << 1 | 1, mid + 1, r);
  pushup(u);
}

void update(int u, int l, int r, int k) {
  if (l <= tr[u].l && tr[u].r <= r) {
    tr[u].add += k;
    tr[u].sum += k * (tr[u].r - tr[u].l + 1);
    return;
  }
  pushdown(u);
  int mid = tr[u].l + tr[u].r >> 1;
  if (l <= mid) update(u << 1, l, r, k);
  if (r > mid) update(u << 1 | 1, l, r, k);
  pushup(u);
}

long query(int u, int l, int r) {
  if (l <= tr[u].l && tr[u].r <= r) return tr[u].sum;
  pushdown(u);
  int mid = tr[u].l + tr[u].r >> 1;
  long res = 0;
  if (l <= mid) res += query(u << 1, l, r);
  if (r > mid) res += query(u << 1 | 1, l, r);
  return res;
}

void update_path(int u, int v, int k) {
  while (top[u] != top[v]) {
    if (dep[top[u]] < dep[top[v]]) swap(u, v);
    // u的重链头更深，并且u重链头在dfs序里下标更小，直接更新u重链头到u这段区间
    update(1, id[top[u]], id[u], k);
    // u跳到重链头上面
    u = fa[top[u]];
  }
  if (dep[u] < dep[v]) swap(u, v);
  update(1, id[v], id[u], k);
}

long query_path(int u, int v) {
  long res = 0;
  while (top[u] != top[v]) {
    if (dep[top[u]] < dep[top[v]]) swap(u, v);
    res += query(1, id[top[u]], id[u]);
    u = fa[top[u]];
  }
  if (dep[u] < dep[v]) swap(u, v);
  res += query(1, id[v], id[u]);
  return res;
}

void update_tree(int u, int k) {
  update(1, id[u], id[u] + sz[u] - 1, k);
}

long query_tree(int u) {
  return query(1, id[u], id[u] + sz[u] - 1);
}

int main() {
  memset(h, -1, sizeof h);
  scanf("%d", &n);
  for (int i = 1; i <= n; i++) scanf("%d", &w[i]);
  for (int i = 1; i < n; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    add(a, b), add(b, a);
  }
  
  dfs1(1, 1, 0);
  dfs2(1, 1);
  
  build(1, 1, n);
  
  scanf("%d", &m);
  while (m--) {
    int t, u, v, k;
    scanf("%d%d", &t, &u);
    if (t == 1) {
      scanf("%d%d", &v, &k);
      update_path(u, v, k);
    } else if (t == 2) {
      scanf("%d", &k);
      update_tree(u, k);
    } else if (t == 3) {
      scanf("%d", &v);
      printf("%ld\n", query_path(u, v));
    } else printf("%ld\n", query_tree(u));
  }
}

预处理时间复杂度$O(n)$，每次询问时间$O({\log }^{2}n)$，空间$O(n)$。