【ACWing】2524. 树链剖分II

题目地址：

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

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

输入格式：
第一行包含一个整数$n$，表示节点个数。
第二行包含$n$个整数，其中第$i$个整数表示$a_i$。
第三行包含$n-1$个整数，其中第$i$个整数表示第$i+1$号节点的父节点编号。
第四行包含一个整数$m$，表示操作次数。
接下来$m$行，每行包含一个操作，格式如题目所述。

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

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

先做树链剖分。操作$2$和操作$4$和谁是树根没关系，可以按照模板https://blog.csdn.net/qq_46105170/article/details/125497358直接操作。最主要的是如何处理换根。原树根为$1$，开个全局变量$r$，存当前的树根。接下来处理询问，假设是询问$u$子树的总权值，分三种情况讨论：
1、如果$u=r$，则就是询问整棵树的总权值；
2、如果$u$和$r$的最近公共祖先不是$u$，那么修改$u$子树与$1$是树根还是$r$是树根没关系，可以当成是对原树的询问；
3、如果$u$和$r$的最近公共祖先是$u$，说明$r$在$u$子树里，那么要求沿着$r$向上走到$u$的下面一步的节点$v$，先询问整棵树的总权值$S$，再询问$v$子树的权值$s$，答案即为$S-s$。

找最近公共祖先也可以用树链剖分做，每次将重链头深的那个节点向上跳，每次跳到重链头的父节点处，直到两个节点重链头相等，此时深度浅的即为最近公共祖先。

求沿着$r$向上走到$u$的下面一步的节点$v$，这一操作也可以用树链剖分。只要$r$的重链头与$u$的重链头不一样就进行循环，如果$r$的重链头的父亲就是$u$，则说明$v$就是$r$重链头；否则就将$r$跳到其重链头的父亲处。退出循环的时候说明$r$和$u$已经在同一条重链上了，那么$v$就是$u$的重儿子。

对子树进行修改的操作与上面类似做。代码如下：

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

const int N = 1e5 + 10;
int n, m;
int h[N], e[N], ne[N], idx, w[N];
int id[N], nw[N], cnt;
int dep[N], son[N], top[N], sz[N], fa[N];
int root;
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++;
}

void dfs1(int u, int depth) {
  dep[u] = depth, sz[u] = 1;
  for (int i = h[u]; ~i; i = ne[i]) {
    int v = e[i];
    dfs1(v, depth + 1);
    sz[u] += sz[v];
    if (sz[son[u]] < sz[v]) son[u] = v;
  }
}

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 == 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], &l = tr[u << 1], &r = tr[u << 1 | 1];
  if (root.add) {
    l.sum += root.add * (l.r - l.l + 1);
    l.add += root.add;
    r.sum += root.add * (r.r - r.l + 1);
    r.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 modify(int u, int l, int r, long 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) modify(u << 1, l, r, k);
  if (r > mid) modify(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 add_path(int u, int v, int k) {
  while (top[u] != top[v]) {
    if (dep[top[u]] < dep[top[v]]) swap(u, v);
    modify(1, id[top[u]], id[u], k);
    u = fa[top[u]];
  }
  if (dep[u] < dep[v]) swap(u, v);
  modify(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;
}

int lca(int u, int v) {
  while (top[u] != top[v]) {
    // 谁的重链头深度深，谁就向上跳
    if (dep[top[u]] < dep[top[v]]) swap(u, v);
    u = fa[top[u]];
  }
  return dep[u] < dep[v] ? u : v;
}

// 输入保证u属于v子树
int get_son(int u, int v) {
  while (top[u] != top[v]) {
    if (fa[top[u]] == v) return top[u];
    u = fa[top[u]];
  }
  return son[v];
}

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

void add_tree(int u, int k) {
  if (u == root) modify(1, 1, n, k);
  else if (lca(u, root) != u) modify(1, id[u], id[u] + sz[u] - 1, k);
  else {
    int x = get_son(root, u);
    modify(1, 1, n, k);
    modify(1, id[x], id[x] + sz[x] - 1, -k);
  }
}

long query_tree(int u) {
  if (u == root) return query(1, 1, n);
  if (lca(u, root) != u) return query(1, id[u], id[u] + sz[u] - 1);
  int x = get_son(root, u);
  return query(1, 1, n) - query(1, id[x], id[x] + sz[x] - 1);
}

int main() {
  memset(h, -1, sizeof h);
  scanf("%d", &n);
  for (int i = 1; i <= n; i++) scanf("%ld", &w[i]);
  for (int i = 2; i <= n; i++) {
    int p;
    scanf("%d", &p);
    add(p, i);
    fa[i] = p;
  }

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

  scanf("%d", &m);
  while (m--) {
    int op, u, v;
    long k;
    scanf("%d%d", &op, &u);
    if (op == 1) {
      root = u;
    } else if (op == 2) {
      scanf("%d%ld", &v, &k);
      add_path(u, v, k);
    } else if (op == 3) {
      scanf("%ld", &k);
      add_tree(u, k);
    } else if (op == 4) {
      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)$。