# Description
P4178 Tree
# Solution
前两天 vp 时碰到一道需要在每个点维护一棵平衡树,同时需要树上启发式合并的题,然鹅发现自己不会写。
无意中发现这道题似乎也可以使用同样的做法。
大概思路是,对于每个点,先继承它的重儿子的平衡树,然后依次枚举所有的轻儿子,先统计答案,再合并平衡树。
统计答案与合并平衡树时,直接遍历轻儿子平衡树内所有的点即可。
具体来说,平衡树中存的值是子树内每个点到当前平衡树的根的点的距离。
假设当前处理到 节点,我们先计算一端为 ,另一端在重儿子子树内的合法点对数,即在重儿子的平衡树中查找小于等于 的点的个数。
然后我们把 点插入到平衡树中。注意到这里平衡树的根从 变成了 ,这意味着需要一个平衡树整体加 的操作,但是整体加还是太吃操作了,我们考虑把后面插入的节点都减去这个值也可以达到同样的效果。
接着枚举所有轻儿子,统计答案时遍历轻儿子的每一个点,取当前重儿子的平衡树中查询小于等于 的点个数即可( 是遍历到的点, 是 所在的子树的根,是 的一个轻儿子)。
询问就处理完了。
合并也是同样的道理,只需要遍历轻儿子,将 插入到重儿子的平衡树中即可。
实际上上面出现的所有距离的公式并不完全准确,只是举个例子,因为我们一直在搜索回溯,每次处理一个重儿子,距离都要整体加上一个数值,具体看代码理解吧。
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
using namespace std;
const int N = 1e6 + 10;
typedef pair<int, int> pii;
int n, k;
vector <pii> G[N];
struct Tree{
int siz[N], val[N], wei[N], ls[N], rs[N];
int tot;
inline void pushup(int x){
siz[x] = siz[ls[x]] + siz[rs[x]] + 1;
}
inline void split(int x, int k, int &a, int &b){
if(!x) return a = b = 0, void();
if(k >= val[x]) a = x, split(rs[x], k, rs[x], b);
else b = x, split(ls[x], k, a, ls[x]);
pushup(x);
}
inline int merge(int x, int y){
if(!x || !y) return x | y;
if(wei[x] <= wei[y]){
rs[x] = merge(rs[x], y);
return pushup(x), x;
}else{
ls[y] = merge(x, ls[y]);
return pushup(y), y;
}
}
inline int newnode(int k){
siz[++tot] = 1, val[tot] = k, wei[tot] = rand();
return tot;
}
inline void ins(int &x, int k){
int a, b; split(x, k, a, b);
x = merge(a, merge(newnode(k), b));
}
inline int qry(int x, int k){
int a, b; split(x, k, a, b);
int res = siz[a];
return x = merge(a, b), res;
}
inline int query(int &p, int q, int du, int dv){
if(!q) return 0;
int tmp = k - du - dv - val[q];
int cnt = qry(p, tmp);
return cnt + query(p, ls[q], du, dv) + query(p, rs[q], du, dv);
}
inline void Union(int &p, int q, int du, int dv){
if(!q) return;
ins(p, val[q] + dv - du);
Union(p, ls[q], du, dv), Union(p, rs[q], du, dv);
}
} t;
int siz[N], root[N], son[N];
int dis[N], sw[N];
int ans;
inline void dfs(int x, int fa){
siz[x] = 1;
for(auto y : G[x]){
if(y.fi == fa) continue;
dfs(y.fi, x), siz[x] += siz[y.fi];
if(siz[y.fi] > siz[son[x]]) son[x] = y.fi, sw[x] = y.se;
}
}
inline void dfs2(int x, int fa, int d){
if(!son[x]){
root[x] = t.newnode(0), dis[x] = d;
return;
}
dfs2(son[x], x, sw[x]);
root[x] = root[son[x]], dis[x] = dis[son[x]];
ans += t.qry(root[x], k - dis[x]);
// x ~ son[x]
// 区间加
t.ins(root[x], -dis[x]); //插入 x
for(auto y : G[x]){
if(y.fi == fa || y.fi == son[x]) continue;
dfs2(y.fi, x, y.se);
ans += t.query(root[x], root[y.fi], dis[x], dis[y.fi]);
t.Union(root[x], root[y.fi], dis[x], dis[y.fi]);
}
dis[x] += d;
}
signed main(){
freopen("test.in", "r", stdin);
freopen("test.out", "w", stdout);
ios :: sync_with_stdio(false);
cin >> n;
for(int i = 1; i < n; ++i){
int u, v, w; cin >> u >> v >> w;
G[u].push_back(pii(v, w)), G[v].push_back(pii(u, w));
}
cin >> k;
dfs(1, 0), dfs2(1, 0, 0);
cout << ans << endl;
return 0;
}