← 返回题目页

A1499 【模板】重链剖分/树链剖分

标程

// A1499 【模板】重链剖分 / 树链剖分
//
// 树链剖分 + 线段树,维护四种操作:
//   1 x y z:x 到 y 路径上所有点加 z
//   2 x y  :x 到 y 路径上所有点的和
//   3 x z  :x 的子树内所有点加 z
//   4 x    :x 的子树内所有点的和
//
// 关键:把树上路径拆成 O(log n) 段连续区间,把子树变成一段连续区间,
// 于是问题变成「区间加、区间求和」的线段树。
// 所有数值都对 P 取模(加法与求和都与取模相容,随时取模不影响结果)。

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int MAXN = 100000 + 5;

int n, m, root, mod;
ll val[MAXN];
vector<int> g[MAXN];
int fa[MAXN], son[MAXN], sz[MAXN], top[MAXN], dfn[MAXN], rid[MAXN], dep[MAXN];
ll seg[MAXN * 4], lazy[MAXN * 4];

/* ------------------------------------------------------------- 线段树 */

void pushUp(int p) {
    seg[p] = (seg[p << 1] + seg[p << 1 | 1]) % mod;
}

void apply(int p, int l, int r, ll v) {
    seg[p] = (seg[p] + v * (r - l + 1)) % mod;
    lazy[p] = (lazy[p] + v) % mod;
}

void pushDown(int p, int l, int r) {
    if (lazy[p] == 0) return;
    int mid = (l + r) >> 1;
    apply(p << 1, l, mid, lazy[p]);
    apply(p << 1 | 1, mid + 1, r, lazy[p]);
    lazy[p] = 0;
}

void addSeg(int p, int l, int r, int ql, int qr, ll v) {
    if (ql <= l && r <= qr) {
        apply(p, l, r, v);
        return;
    }
    pushDown(p, l, r);
    int mid = (l + r) >> 1;
    if (ql <= mid) addSeg(p << 1, l, mid, ql, qr, v);
    if (qr > mid) addSeg(p << 1 | 1, mid + 1, r, ql, qr, v);
    pushUp(p);
}

ll sumSeg(int p, int l, int r, int ql, int qr) {
    if (ql <= l && r <= qr) return seg[p];
    pushDown(p, l, r);
    int mid = (l + r) >> 1;
    ll res = 0;
    if (ql <= mid) res = (res + sumSeg(p << 1, l, mid, ql, qr)) % mod;
    if (qr > mid) res = (res + sumSeg(p << 1 | 1, mid + 1, r, ql, qr)) % mod;
    return res;
}

void addRange(int l, int r, ll v) {
    if (l > r) return;
    addSeg(1, 1, n, l, r, v);
}

ll sumRange(int l, int r) {
    if (l > r) return 0;
    return sumSeg(1, 1, n, l, r);
}

/* --------------------------------------------------------- 树链剖分 */

// 第一遍:算父亲、深度、子树大小、重儿子(子树最大的那个孩子)
void buildSizes(int r) {
    vector<int> order;
    order.push_back(r);
    fa[r] = 0;
    dep[r] = 1;
    for (size_t i = 0; i < order.size(); i++) {
        int u = order[i];
        sz[u] = 1;
        son[u] = 0;
        for (int v : g[u]) {
            if (v == fa[u]) continue;
            fa[v] = u;
            dep[v] = dep[u] + 1;
            order.push_back(v);
        }
    }
    for (int i = (int)order.size() - 1; i >= 0; i--) {
        int u = order[i];
        for (int v : g[u]) {
            if (v == fa[u]) continue;
            sz[u] += sz[v];
            if (!son[u] || sz[v] > sz[son[u]]) son[u] = v;
        }
    }
}

// 第二遍:重儿子优先的 DFS 序,dfn[x] 是 x 在序列里的位置
void buildDfn(int r) {
    vector<pair<int, int>> stk;
    stk.push_back({r, r});
    int timer = 0;
    while (!stk.empty()) {
        int u = stk.back().first, tp = stk.back().second;
        stk.pop_back();
        for (int x = u; x; x = son[x]) {
            top[x] = tp;
            dfn[x] = ++timer;
            rid[timer] = x;
            vector<int> lights;
            for (int v : g[x]) {
                if (v == fa[x] || v == son[x]) continue;
                lights.push_back(v);
            }
            for (int i = (int)lights.size() - 1; i >= 0; i--) stk.push_back({lights[i], lights[i]});
        }
    }
}

// 子树 x 对应的连续区间
int subtreeL(int x) { return dfn[x]; }
int subtreeR(int x) { return dfn[x] + sz[x] - 1; }

// 路径 x -> y 拆成若干区间,逐段加
void pathAdd(int x, int y, ll v) {
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        addRange(dfn[top[x]], dfn[x], v);
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    addRange(dfn[x], dfn[y], v);
}

// 路径 x -> y 拆成若干区间,逐段求和
ll pathSum(int x, int y) {
    ll res = 0;
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        res = (res + sumRange(dfn[top[x]], dfn[x])) % mod;
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    res = (res + sumRange(dfn[x], dfn[y])) % mod;
    return res;
}

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

    cin >> n >> m >> root >> mod;
    for (int i = 1; i <= n; i++) {
        cin >> val[i];
        val[i] %= mod;
    }
    for (int i = 1; i < n; i++) {
        int x, y;
        cin >> x >> y;
        g[x].push_back(y);
        g[y].push_back(x);
    }

    buildSizes(root);
    buildDfn(root);

    // 先把所有点塞进线段树(按 dfn 序)
    {
        // 建树
        vector<ll> a(n + 1);
        for (int i = 1; i <= n; i++) a[i] = val[rid[i]];
        function<void(int, int, int)> build = [&](int p, int l, int r) {
            if (l == r) {
                seg[p] = a[l] % mod;
                return;
            }
            int mid = (l + r) >> 1;
            build(p << 1, l, mid);
            build(p << 1 | 1, mid + 1, r);
            pushUp(p);
        };
        build(1, 1, n);
    }

    while (m--) {
        int op;
        cin >> op;
        if (op == 1) {
            int x, y;
            ll z;
            cin >> x >> y >> z;
            pathAdd(x, y, z % mod);
        } else if (op == 2) {
            int x, y;
            cin >> x >> y;
            cout << pathSum(x, y) << '\n';
        } else if (op == 3) {
            int x;
            ll z;
            cin >> x >> z;
            addRange(subtreeL(x), subtreeR(x), z % mod);
        } else if (op == 4) {
            int x;
            cin >> x;
            cout << sumRange(subtreeL(x), subtreeR(x)) << '\n';
        }
    }
    return 0;
}

动画演示