【洛谷】P3384 【模板】轻重链剖分/树链剖分

题目地址：

https://www.luogu.com.cn/problem/P3384

题目描述：
如题，已知一棵包含$N$个结点的树（连通且无环），每个节点上包含一个数值，需要支持以下操作：
1 x y z，表示将树从$x$到$y$结点最短路径上所有节点的值都加上$z$。
2 x y，表示求树从$x$到$y$结点最短路径上所有节点的值之和。
3 x z，表示将以$x$为根节点的子树内所有节点值都加上$z$。
4 x表示求以$x$为根节点的子树内所有节点值之和

输入格式：
第一行包含$4$个正整数$N,M,R,P$，分别表示树的结点个数、操作个数、根节点序号和取模数（即所有的输出结果均对此取模）。
接下来一行包含$N$个非负整数，分别依次表示各个节点上初始的数值。
接下来$N-1$行每行包含两个整数$x,y$，表示点$x$和点$y$之间连有一条边（保证无环且连通）。
接下来$M$行每行包含若干个正整数，每行表示一个操作。

输出格式：
输出包含若干行，分别依次表示每个操作$2$或操作$4$所得的结果（对$P$取模）。

数据范围：
对于 $30\%$ 的数据： $1\le N\le 10$，$1\le M\le 10$；
对于 $70\%$ 的数据： $1\le N\le {10}^3$，$1\le M\le {10}^3$；
对于 $100\%$ 的数据： $1\le N\le {10}^5$，$1\le M\le {10}^5$，$1\le R\le N$，$1\le P\le 2^{31}-1$。

思路是树链剖分。参考：https://blog.csdn.net/qq_46105170/article/details/125497358。代码如下：

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

const int N = 1e5 + 10, M = N * 2;
int n, m, rt, P;
int h[N], e[M], ne[M], idx;
long w[N], nw[N];
int id[N], cnt;
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++;
}

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;
  }
}

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], &l = tr[u << 1], &r = tr[u << 1 | 1];
  if (root.add) {
    l.sum = (l.sum + root.add * (l.r - l.l + 1)) % P;
    l.add = (l.add + root.add) % P;
    r.sum = (r.sum + root.add * (r.r - r.l + 1)) % P;
    r.add = (r.add + root.add) % P;
    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, long k) {
  if (l <= tr[u].l && tr[u].r <= r) {
    tr[u].add = (tr[u].add + k) % P;
    tr[u].sum = (tr[u].sum + k * (tr[u].r - tr[u].l + 1)) % P;
    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 = (res + query(u << 1, l, r)) % P;
  if (r > mid) res = (res + query(u << 1 | 1, l, r)) % P;
  return res;
}

void update_path(int u, int v, long k) {
  while (top[u] != top[v]) {
    if (dep[top[u]] < dep[top[v]]) swap(u, v);
    update(1, id[top[u]], id[u], k);
    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 = (res + query(1, id[top[u]], id[u])) % P;
    u = fa[top[u]];
  }
  if (dep[u] < dep[v]) swap(u, v);
  res = (res + query(1, id[v], id[u])) % P;
  return res;
}

int main() {
  memset(h, -1, sizeof h);
  scanf("%d%d%d%d", &n, &m, &rt, &P);
  for (int i = 1; i <= n; i++) scanf("%ld", &w[i]);
  for (int i = 1; i < n; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    add(a, b), add(b, a);
  }

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

  while (m--) {
    int t, u, v;
    long k;
    scanf("%d%d", &t, &u);
    if (t == 1) {
      scanf("%d%ld", &v, &k);
      update_path(u, v, k);
    } else if (t == 2) {
      scanf("%d", &v);
      printf("%ld\n", query_path(u, v));
    } else if (t == 3) {
      scanf("%ld", &k);
      update(1, id[u], id[u] + sz[u] - 1, k);
    } else printf("%ld\n", query(1, id[u], id[u] + sz[u] - 1));
  }
}

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