【题解】HDU-4918-Query on the subtree

题目

http://acm.hdu.edu.cn/showproblem.php?pid=4918

题意

给一棵有点权的树,要求支持两种操作:

  • 修改一个点的点权。
  • 查询到 $u$ 距离不超过 $d$ 的点的点权和。

题解

  • 考虑点分治形成的重心树,点 $u$ 至多出现在 $\log n$ 个重心的子树上。
  • 假设点 $u$ 在重心 $rt$ 的子树上,且在 $rt$ 的儿子 $subrt$ 的子树上。设 $u$ 在 $rt$ 中的深度为 $dep$,$rt.sum[d]$ 是 $rt$ 的子树中深度不超过 $d$ 的点的点权和,那么 $u$ 在 $rt$ 中的贡献是 $rt.sum[d-dep] - subrt.sum[d-dep]$。
  • 利用树状数组就能支持查询和修改的操作,但是需要记录每个点所有的 $rt$ 和 $subrt$ 以及对应的深度。

代码

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
#include <bits/stdc++.h>
using namespace std;
using LL = long long;
#define FOR(i, x, y) for (decay<decltype(y)>::type i = (x), _##i = (y); i < _##i; ++i)
#define FORD(i, x, y) for (decay<decltype(x)>::type i = (x), _##i = (y); i > _##i; --i)
#ifdef zerol
#define dbg(args...) do { cout << "\033[32;1m" << #args << " -> "; err(args); } while (0)
#else
#define dbg(...)
#endif
void err() { cout << "\033[39;0m" << endl; }
template<template<typename...> class T, typename t, typename... Args>
void err(T<t> a, Args... args) { for (auto x: a) cout << x << ' '; err(args...); }
template<typename T, typename... Args>
void err(T a, Args... args) { cout << a << ' '; err(args...); }
// -----------------------------------------------------------------------------
const int maxn = 1E5 + 100, INF = 1E9;
int w[maxn];
vector<int> G[maxn];
bool vis[maxn];
int sz[maxn];

struct bit {
int n;
vector<int> c;
bit* init(int _n) { n = _n + 2; c.resize(n); for (int& x: c) x = 0; return this; }
inline int lowbit(int x) { return x & -x; }
void add(int p, int v) {
assert(p >= 0 && p <= n - 2);
for (int i = p + 1; i < n; i += lowbit(i))
c[i] += v;
}
int sum(int p) {
int ret = 0;
for (int i = min(p + 1, n - 1); i > 0; i -= lowbit(i))
ret += c[i];
return ret;
}
} pool[maxn << 1], *pit;

struct P {
bit *rt, *subrt;
int dep;
};
vector<P> R[maxn];

void get_sz(int u, int fa) {
sz[u] = 1;
for (int& v: G[u]) {
if (v == fa || vis[v]) continue;
get_sz(v, u);
sz[u] += sz[v];
}
}
void find_rt(int u, int fa, int s, int& m, int& rt) {
int t = s - sz[u];
for (int& v: G[u]) {
if (v == fa || vis[v]) continue;
find_rt(v, u, s, m, rt);
t = max(t, sz[v]);
}
if (t < m) { m = t; rt = u; }
}

int dep[maxn], md[maxn];
void get_dep(int u, int fa, int d) {
dep[u] = d; md[u] = 0;
for (int& v: G[u]) {
if (v == fa || vis[v]) continue;
get_dep(v, u, d + 1);
md[u] = max(md[u], md[v] + 1);
}
}

void go(int u, int fa, bit* rt, bit* subrt) {
R[u].push_back({rt, subrt, dep[u]});
for (int& v: G[u]) if (v != fa && !vis[v]) go(v, u, rt, subrt);
}

void dfs(int u) {
int tmp = INF;
get_sz(u, -1);
find_rt(u, -1, sz[u], tmp, u);
vis[u] = 1;
get_dep(u, -1, 0);
bit* rt = (pit++)->init(md[u]);
R[u].push_back({rt, nullptr, 0});
for (int& v: G[u]) {
if (vis[v]) continue;
go(v, u, rt, (pit++)->init(md[v] + 1));
}
for (int& v: G[u]) if (!vis[v]) dfs(v);
}

void init(int);
int main() {
int n, Q;
while (cin >> n >> Q) {
init(n);
FOR (i, 1, n + 1) scanf("%d", &w[i]);
FOR (_, 1, n) {
int u, v; scanf("%d%d", &u, &v);
G[u].push_back(v); G[v].push_back(u);
}
dfs(1);
FOR (i, 1, n + 1)
for (P& p: R[i]) {
p.rt->add(p.dep, w[i]);
if (p.subrt) p.subrt->add(p.dep, w[i]);
}
char s[10];
while (Q--) {
int u, x;
scanf("%s%d%d", s, &u, &x);
if (s[0] == '!') {
for (P& p: R[u]) {
p.rt->add(p.dep, x - w[u]);
if (p.subrt) p.subrt->add(p.dep, x - w[u]);
}
w[u] = x;
} else {
int ans = 0;
for (P& p: R[u]) {
ans += p.rt->sum(x - p.dep);
if (p.subrt) ans -= p.subrt->sum(x - p.dep);
}
printf("%d\n", ans);
}
}
}
}
void init(int n) {
FOR (i, 1, n + 1) G[i].clear(), R[i].clear();
pit = pool;
memset(vis, 0, sizeof vis);
}