[SHOI2012] 魔法树

题意

给一棵根为 0 的树，初始时每个点的果子数都是 0。

有两类操作：

A u v d
把路径 u -> v 上所有节点的果子数都加上 d
Q u
询问当前以 u 为根的整棵子树里，一共有多少个果子

样例树

样例树是一条链：

graph G {
  0 -- 1;
  1 -- 2;
  2 -- 3;
}

执行一次 A 1 3 1 后：

点 1, 2, 3 各加 1

所以：

子树 0 的和是 3
子树 1 的和是 3
子树 2 的和是 2

思路

先看一个最直接的小数据暴力：

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116// brute.cpp：真的把路径上的点都加上 d，再暴力统计子树和。
// 复杂度较高，只适合小数据理解和对拍。
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 25;

int n;
vector<int> g[MAXN];
int parent_arr[MAXN];
int depth_arr[MAXN];
long long val[MAXN];

void build_parent_depth() {
    vector<int> st;
    st.push_back(0);
    parent_arr[0] = 0;
    depth_arr[0] = 0;

    while (!st.empty()) {
        int u = st.back();
        st.pop_back();

        for (size_t i = 0; i < g[u].size(); i++) {
            int v = g[u][i];
            if (v == parent_arr[u]) {
                continue;
            }
            parent_arr[v] = u;
            depth_arr[v] = depth_arr[u] + 1;
            st.push_back(v);
        }
    }
}

int lca(int a, int b) {
    while (depth_arr[a] > depth_arr[b]) {
        a = parent_arr[a];
    }
    while (depth_arr[b] > depth_arr[a]) {
        b = parent_arr[b];
    }
    while (a != b) {
        a = parent_arr[a];
        b = parent_arr[b];
    }
    return a;
}

void add_path(int u, int v, long long d) {
    int p = lca(u, v);

    while (u != p) {
        val[u] += d;
        u = parent_arr[u];
    }
    while (v != p) {
        val[v] += d;
        v = parent_arr[v];
    }
    val[p] += d;
}

long long dfs_sum(int u, int fa) {
    long long res = val[u];
    for (size_t i = 0; i < g[u].size(); i++) {
        int v = g[u][i];
        if (v == fa) {
            continue;
        }
        res += dfs_sum(v, u);
    }
    return res;
}

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

    cin >> n;
    for (int i = 0; i < n; i++) {
        g[i].clear();
        parent_arr[i] = 0;
        depth_arr[i] = 0;
        val[i] = 0;
    }

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

    build_parent_depth();

    int q;
    cin >> q;
    while (q--) {
        char op;
        cin >> op;
        if (op == 'A') {
            int u, v;
            long long d;
            cin >> u >> v >> d;
            add_path(u, v, d);
        }
        else {
            int u;
            cin >> u;
            cout << dfs_sum(u, parent_arr[u]) << '\n';
        }
    }

    return 0;
}

暴力做法就是：

真正把路径 u -> v 上的每个点都加上 d
查询时再暴力把子树里的点全部加起来

这个做法最贴近题意，但显然扛不住大数据。

这题的关键不是树剖本身，而是先把“路径加”换成“树上差分”。

对于一次路径加 u, v, d，设 p = lca(u, v)，做：

diff[u] += d
diff[v] += d
diff[p] -= d
diff[parent(p)] -= d（如果 p 不是根）

这样做完以后，某个点 x 当前真正的果子数，等于：

diff 在 x 子树里的总和

也就是：

value[x] = sum(diff[y]), y 在 subtree(x)

现在题目要求的是：

subtree(u) 里所有 value[x] 的总和

把这个式子再展开一层：

sum_{x in subtree(u)} value[x]

等于

sum_{y in subtree(u)} diff[y] * (depth[y] - depth[u] + 1)

原因是：

一个点 y 的差分值，会贡献给它在 subtree(u) 里的所有祖先
这样的祖先数量正好是 depth[y] - depth[u] + 1

于是查询可以改写成：

sum(diff * (depth+1)) - depth[u] * sum(diff)

这就很适合用 DFS 序来维护：

子树在 DFS 序上是连续区间
用一棵树状数组维护 diff
再用另一棵树状数组维护 diff * (depth+1)

路径加时只做 4 次单点修改。
查询子树和时，只要在 [tin[u], tout[u]] 上查这两个区间和，再套公式即可。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199#include <bits/stdc++.h>
using namespace std;

const int MAXN = 100000 + 5;
const int LOG = 18;

struct Fenwick {
    int n;
    long long bit[MAXN];

    void init(int n_) {
        n = n_;
        for (int i = 1; i <= n; i++) {
            bit[i] = 0;
        }
    }

    void add(int x, long long delta) {
        for (int i = x; i <= n; i += i & -i) {
            bit[i] += delta;
        }
    }

    long long sum_prefix(int x) const {
        long long res = 0;
        for (int i = x; i > 0; i -= i & -i) {
            res += bit[i];
        }
        return res;
    }

    long long range_sum(int l, int r) const {
        if (l > r) {
            return 0;
        }
        return sum_prefix(r) - sum_prefix(l - 1);
    }
};

int n;
vector<int> g[MAXN];
int parent_arr[MAXN];
int depth_arr[MAXN];
int up[MAXN][LOG];
int tin[MAXN], tout[MAXN], dfs_clock;
int iter_idx[MAXN];

Fenwick bit_cnt;   // 维护树上差分值 diff
Fenwick bit_depth; // 维护 diff * (depth + 1)

void init_tree(int n) {
    dfs_clock = 0;
    for (int i = 0; i < n; i++) {
        g[i].clear();
        parent_arr[i] = 0;
        depth_arr[i] = 0;
        tin[i] = 0;
        tout[i] = 0;
        iter_idx[i] = 0;
        for (int j = 0; j < LOG; j++) {
            up[i][j] = 0;
        }
    }
}

void add_edge(int u, int v) {
    g[u].push_back(v);
    g[v].push_back(u);
}

void build_lca_and_dfn() {
    vector<int> st;
    st.push_back(0);
    parent_arr[0] = 0;
    depth_arr[0] = 0;
    tin[0] = ++dfs_clock;

    while (!st.empty()) {
        int u = st.back();

        if (iter_idx[u] < (int)g[u].size()) {
            int v = g[u][iter_idx[u]++];
            if (v == parent_arr[u]) {
                continue;
            }

            parent_arr[v] = u;
            depth_arr[v] = depth_arr[u] + 1;
            up[v][0] = u;
            for (int j = 1; j < LOG; j++) {
                up[v][j] = up[up[v][j - 1]][j - 1];
            }

            tin[v] = ++dfs_clock;
            st.push_back(v);
        }
        else {
            tout[u] = dfs_clock;
            st.pop_back();
        }
    }
}

int kth_ancestor(int u, int k) {
    for (int j = 0; j < LOG; j++) {
        if (k & (1 << j)) {
            u = up[u][j];
        }
    }
    return u;
}

int lca(int a, int b) {
    if (depth_arr[a] < depth_arr[b]) {
        swap(a, b);
    }

    a = kth_ancestor(a, depth_arr[a] - depth_arr[b]);
    if (a == b) {
        return a;
    }

    for (int j = LOG - 1; j >= 0; j--) {
        if (up[a][j] != up[b][j]) {
            a = up[a][j];
            b = up[b][j];
        }
    }

    return up[a][0];
}

void add_diff(int u, long long delta) {
    int pos = tin[u];
    bit_cnt.add(pos, delta);
    bit_depth.add(pos, delta * (depth_arr[u] + 1LL));
}

// 对路径 u-v 上所有点加 d。
// 使用树上差分：diff[u] += d, diff[v] += d, diff[lca] -= d, diff[parent(lca)] -= d。
void add_path(int u, int v, long long d) {
    int p = lca(u, v);

    add_diff(u, d);
    add_diff(v, d);
    add_diff(p, -d);
    if (p != 0) {
        add_diff(parent_arr[p], -d);
    }
}

// 当前子树和：
// sum_{x in subtree(u)} value[x]
// = sum_{y in subtree(u)} diff[y] * (depth[y] - depth[u] + 1)
// = sum(diff * (depth+1)) - depth[u] * sum(diff)
long long query_subtree(int u) {
    long long s1 = bit_cnt.range_sum(tin[u], tout[u]);
    long long s2 = bit_depth.range_sum(tin[u], tout[u]);
    return s2 - 1LL * depth_arr[u] * s1;
}

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

    cin >> n;
    init_tree(n);

    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        add_edge(u, v);
    }

    build_lca_and_dfn();
    bit_cnt.init(n);
    bit_depth.init(n);

    int q;
    cin >> q;
    while (q--) {
        char op;
        cin >> op;

        if (op == 'A') {
            int u, v;
            long long d;
            cin >> u >> v >> d;
            add_path(u, v, d);
        }
        else {
            int u;
            cin >> u;
            cout << query_subtree(u) << '\n';
        }
    }

    return 0;
}

复杂度

预处理：

DFS 序和倍增 LCA： $O(n log n)$

每次操作：

路径加：1 次 LCA + 4 次树状数组单点修改， $O(log n)$
子树查询：2 次树状数组区间求和， $O(log n)$

空间复杂度：

$O(n log n)$

总结

这题最值得记住的不是“树状数组”本身，而是这一步转化：

路径加先做树上差分
子树和再把 value[x] 展开成 diff 的加权和

一旦写成：

sum(diff * (depth+1)) - depth[u] * sum(diff)

整题就从“树上路径更新 + 子树查询”变成了：

LCA
DFS 序
两棵树状数组

的经典组合。

一图流解析

这张图把本题的建模、关键转移、实现检查和训练方法压缩到一页，适合读完正文后复盘。

目录

题意

样例树

思路

代码

复杂度

总结

一图流解析