重链剖分/树链剖分模板

洛谷模板题AC通过

#include <bits/stdc++.h>

using namespace std;

using i64 = long long;

int P;

template<class Node>
struct SegmentTree {
#define lc u<<1
#define rc u<<1|1
    const int n, N;
    vector<Node> tr;

    SegmentTree(): n(0) {}
    SegmentTree(int n_): n(n_), N(n * 4 + 10) {
        tr.reserve(N);
        tr.resize(N);
    }
    SegmentTree(vector<int> init) : SegmentTree(init.size() - 1) {
        function<void(int, int, int)> build = [&](int u, int l, int r) {
            tr[u].l = l, tr[u].r = r;
            init_lazy(tr[u]);
            if (l == r) {
                tr[u] = {l, r, 0, init[l]};
                return ;
            }
            i64 mid = (l + r) >> 1;
            build(lc, l, mid);
            build(rc, mid + 1, r);
            pushup(tr[u], tr[lc], tr[rc]);
        };
        build(1, 1, n);
    }

    void cal_lazy(Node & father, Node & ch) {
        i64 b = father.add % P;
        ch.sum += b * (ch.r - ch.l + 1) % P;
        ch.sum %= P;
    }

    void tag_union(Node& father, Node& ch) {
        i64 b = father.add % P;
        ch.add += b;
        ch.add %= P;
    }

    void init_lazy(Node& u) {
        u.add = 0;
    }

    void pushdown(i64 u) {
        if (tr[u].add != 0) {
            cal_lazy(tr[u], tr[lc]);
            cal_lazy(tr[u], tr[rc]);
            tag_union(tr[u], tr[lc]);
            tag_union(tr[u], tr[rc]);
            init_lazy(tr[u]);
        }
    }

    void pushup(Node& U, Node& L, Node& R) { //上传
        U.l = L.l, U.r = R.r;
        U.sum = (L.sum + R.sum) % P;
    }

    void modify(int u, int l, int r, int k) {
        if (tr[u].l >= l && tr[u].r <= r) {
            (tr[u].add += k) %= P;
            (tr[u].sum += 1LL * k * (tr[u].r - tr[u].l + 1) % P) %= P;
            return ;
        }
        pushdown(u);
        int mid = (tr[u].l + tr[u].r) >> 1;
        if (l <= mid)
            modify(lc, l, r, k);
        if (r > mid)
            modify(rc, l, r, k);
        pushup(tr[u], tr[lc], tr[rc]);
    }

    Node query(int u, int l, int r) { //区查
        if (l <= tr[u].l && tr[u].r <= r)
            return tr[u];
        i64 mid = tr[u].l + tr[u].r >> 1;
        pushdown(u);
        if (r <= mid)
            return query(lc, l, r);
        if (l > mid)
            return query(rc, l, r);
        Node U;
        Node L = query(lc, l, r), R = query(rc, l, r);
        pushup(U, L, R);
        return U;
    }
};

struct Node { //线段树定义
    int l, r, add;
    i64 sum;
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m, r;
    cin >> n >> m >> r >> P;

    vector<int> a(n + 1);
    for (int i = 1; i <= n; i += 1) {
        cin >> a[i];
        a[i] %= P;
    }

    vector e(n + 1, vector<int>());
    for (int i = 1; i < n; i += 1) {
        int u, v;
        cin >> u >> v;
        e[u].emplace_back(v);
        e[v].emplace_back(u);
    }

    vector<int> fa(n + 1), dep(n + 1), son(n + 1), siz(n + 1);
    auto dfs1 = [&](auto && self, int u, int father)->void{
        siz[u] = 1;
        fa[u] = father;
        dep[u] = dep[father] + 1;

        int MaxSon = -1;
        for (auto &v : e[u]) {
            if (v == father) continue;
            self(self, v, u);

            siz[u] += siz[v];
            if (siz[v] > MaxSon) {
                MaxSon = siz[v];
                son[u] = v;
            }
        }
    };
    dfs1(dfs1, r, 0);

    int cnt = 0;
    vector<int> top(n + 1), val(n + 1), dfn(n + 1);
    auto dfs2 = [&](auto && self, int u, int father)->void{
        top[u] = father;
        dfn[u] = ++cnt;
        val[cnt] = a[u];

        if (!son[u]) return;
        self(self, son[u], father);

        for (auto &v : e[u]) {
            if (!dfn[v]) {
                self(self, v, v);
            }
        }
    };
    dfs2(dfs2, r, r);

    SegmentTree<Node> S(val);

    auto modify_tree = [&](int x, int y, int v)->void{
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            S.modify(1, dfn[top[x]], dfn[x], v);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        S.modify(1, dfn[x], dfn[y], v);
    };
    auto query_tree = [&](int x, int y)->int{
        int ans = 0;
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            ans = (1LL * ans + S.query(1, dfn[top[x]], dfn[x]).sum) % P;
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        ans = (1LL * ans + S.query(1, dfn[x], dfn[y]).sum) % P;
        return ans;
    };

    while (m--) {
        int op;
        cin >> op;

        if (op == 1) {
            int x, y, z;
            cin >> x >> y >> z;
            modify_tree(x, y, z % P);
        } else if (op == 2) {
            int x, y;
            cin >> x >> y;
            cout << query_tree(x, y) << "\n";
        } else if (op == 3) {
            int x, z;
            cin >> x >> z;
            S.modify(1, dfn[x], dfn[x] + siz[x] - 1, z);
        } else {
            int x;
            cin >> x;
            cout << S.query(1, dfn[x], dfn[x] + siz[x] - 1).sum << "\n";
        }
    }

    return 0;
}

dfs树链剖分

    vector<int> fa(n + 1), dep(n + 1), son(n + 1), siz(n + 1);
    auto dfs1 = [&](auto && self, int u, int father)->void{
        siz[u] = 1;
        fa[u] = father;
        dep[u] = dep[father] + 1;

        int MaxSon = -1;
        for (auto &v : e[u]) {
            if (v == father) continue;
            self(self, v, u);

            siz[u] += siz[v];
            if (siz[v] > MaxSon) {
                MaxSon = siz[v];
                son[u] = v;
            }
        }
    };
    dfs1(dfs1, r, 0);

    int cnt = 0;
    vector<int> top(n + 1), val(n + 1), dfn(n + 1);
    auto dfs2 = [&](auto && self, int u, int father)->void{
        top[u] = father;
        dfn[u] = ++cnt;
        val[cnt] = a[u];

        if (!son[u]) return;
        self(self, son[u], father);

        for (auto &v : e[u]) {
            if (!dfn[v]) {
                self(self, v, v);
            }
        }
    };
    dfs2(dfs2, r, r);

线段树维护树上两点操作

    SegmentTree<Node> S(val);

    auto modify_tree = [&](int x, int y, int v)->void{
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            S.modify(1, dfn[top[x]], dfn[x], v);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        S.modify(1, dfn[x], dfn[y], v);
    };
    auto query_tree = [&](int x, int y)->int{
        int ans = 0;
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            ans = (1LL * ans + S.query(1, dfn[top[x]], dfn[x]).sum) % P;
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        ans = (1LL * ans + S.query(1, dfn[x], dfn[y]).sum) % P;
        return ans;
    };

树剖求Lca

    auto lca = [&](int x, int y)->int {
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        return x;
    };

树剖换根，当前根是y，求以x为根的包含y的直接儿子（x的子树要包含y）

    auto get_son = [&](int x, int y)->int {
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            if (fa[top[x]] == y) return top[x];
            x = fa[top[x]];
        }

        if (dep[x] > dep[y]) swap(x, y);
        return son[x];
    };

边权转点权，写法不同

    vector<int> fa(n + 1), dep(n + 1), son(n + 1), siz(n + 1), a(n + 1);
    auto dfs1 = [&](auto && self, int u, int father)->void{
        siz[u] = 1;
        fa[u] = father;
        dep[u] = dep[father] + 1;

        int MaxSon = -1;
        for (auto &[v, w] : e[u]) {
            if (v == father) continue;

            a[v] = w;//边权转点权
            self(self, v, u);

            siz[u] += siz[v];
            if (siz[v] > MaxSon) {
                MaxSon = siz[v];
                son[u] = v;
            }
        }
    };
    dfs1(dfs1, 1, 0);

    int cnt = 0;
    vector<int> top(n + 1), val(n + 1), dfn(n + 1);
    auto dfs2 = [&](auto && self, int u, int father)->void{
        top[u] = father;
        dfn[u] = ++cnt;
        val[cnt] = a[u];

        if (!son[u]) return;
        self(self, son[u], father);

        for (auto &[v, _] : e[u]) {
            if (!dfn[v]) {
                self(self, v, v);
            }
        }
    };
    dfs2(dfs2, 1, 1);

    SegmentTree<Node> S(val);

    auto modify_tree = [&](int x, int y, int v)->void{
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            S.modify(1, dfn[top[x]], dfn[x], v);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        if (x != y) {
            S.modify(1, dfn[x] + 1, dfn[y], v);
        }
    };
    auto modify_tree_add = [&](int x, int y, int v)->void{
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            S.modify_add(1, dfn[top[x]], dfn[x], v);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        if (x != y) {
            S.modify_add(1, dfn[x] + 1, dfn[y], v);
        }
    };
    auto query_tree = [&](int x, int y)->int{
        int ans = 0;
        while (top[x] != top[y]) {
            if (dep[top[x]] < dep[top[y]]) swap(x, y);
            ans = max(ans, S.query(1, dfn[top[x]], dfn[x]).Max);
            x = fa[top[x]];
        }
        if (dep[x] > dep[y]) swap(x, y);
        if (x != y) {
            ans = max(ans, S.query(1, dfn[x] + 1, dfn[y]).Max);
        }
        return ans;
    };

参考文章：
[1].【学习笔记】树链剖分系列一 ——重链剖分
[2].【学习笔记】树链剖分系列二——LCA
[3].【学习笔记】树链剖分系列三——边权转点权
[4].【学习笔记】树链剖分系列四——换根操作
[5].换根树剖 学习笔记