SPOJ COT Count on a tree 树上主席树 + LCA

前置技能

可持久化线段树、LCA

题意

给一棵树,求(u,v)的最短路径之间第k小的值。

思路

看到第k小值就是到可以用主席树维护了。 但是一般的主席树是维护一段数列,求区间内的第k值。 不过没关系,我们可以在遍历树的同时建主席树。

那么怎么能求第k值呢?

我们知道主席树也有前缀和的思想,所以我们需要一个媒介点,通过这个点达到前缀和的目的。

与树上两个点有联系的点,那就是LCA–最近公共祖先。

所以我们可以在求LCA_DFS(u, pre)的同时建主席树,同时有需要建的根节点和前一版本的根节点。

然后就是查询部分,根据下图:

答案为: $query(root[l], \;root[r],\; root[lca],\; root[fa[lca][0]],\; 1,\; n,\;k)$

在递归过程中,比较hjt[hjt[ql].l].val + hjt[hjt[qr].l].val - hjt[hjt[lca].l].val - hjt[hjt[falca].l].val与k的大小来判断要进入左子树还是右子树。

Code

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
##include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pdd;

##define INF 0x3f3f3f3f
##define lowbit(x) x & (-x)
##define mem(a, b) memset(a , b , sizeof(a))
##define FOR(i, x, n) for(int i = x;i <= n; i++)

// const ll mod = 998244353;
// const ll mod = 1e9 + 7;
// const double eps = 1e-6;
// const double PI = acos(-1);
// const double R = 0.57721566490153286060651209;

const int N = 2e5 + 10;

int n;

struct Edge {
int u, v, next;
}e[N << 1];
int idx, head[N << 1];
int fa[N][55], depth[N], lg[N];

struct Node {
int l, r;
int val;
}hjt[N * 40];

int a[N], root[N], cnt;
vector<int> vec;

inline void add(int u, int v) {
e[++idx].v = v;
e[idx].next = head[u];
head[u] = idx;
}

inline int getid(int x) { return lower_bound(vec.begin(), vec.end(), x) - vec.begin() + 1; }

void insert(int pre, int &now, int l, int r, int p) {
hjt[now = ++cnt] = hjt[pre];
hjt[now].val++;
if(l == r) return ;
int m = (l + r) / 2;
if(p <= m) insert(hjt[pre].l, hjt[now].l, l, m, p);
else insert(hjt[pre].r, hjt[now].r, m + 1, r, p);
}

int query(int ql, int qr, int lca, int falca, int l, int r, int k) {
if(l == r) return l;
int m = (l + r) / 2;
int tmp = hjt[hjt[ql].l].val + hjt[hjt[qr].l].val - hjt[hjt[lca].l].val - hjt[hjt[falca].l].val;
if(k <= tmp) return query(hjt[ql].l, hjt[qr].l, hjt[lca].l, hjt[falca].l, l, m, k);
else return query(hjt[ql].r, hjt[qr].r, hjt[lca].r, hjt[falca].r, m + 1, r, k - tmp);
}

int LCA(int u1, int u2) {
if(depth[u1] < depth[u2]) swap(u1, u2);
while(depth[u1] > depth[u2])
u1 = fa[u1][lg[depth[u1] - depth[u2]] - 1];
if(u1 == u2) return u1;

for(int i = lg[depth[u1]] - 1;i >= 0; i--) {
if(fa[u1][i] != fa[u2][i]) {
u1 = fa[u1][i];
u2 = fa[u2][i];
}
}
return fa[u1][0];
}

void LCA_DFS(int u, int pre) {
fa[u][0] = pre;
depth[u] = depth[pre] + 1;
insert(root[pre], root[u], 1, n, getid(a[u])); // 沿路径建主席树(前缀和)
for(int i = 1;i <= lg[depth[u]]; i++) {
fa[u][i] = fa[fa[u][i - 1]][i - 1];
}
for(int i = head[u]; i; i = e[i].next) {
if(e[i].v != pre)
LCA_DFS(e[i].v, u);
}
}

void solve()
{
int m;
cin >> n >> m;
for(int i = 1;i <= n; i++) {
cin >> a[i]; vec.push_back(a[i]);
}
sort(vec.begin(), vec.end());
vec.erase(unique(vec.begin(), vec.end()), vec.end());
for(int i = 1;i < n; i++) {
int u, v;
cin >> u >> v;
add(u, v);
add(v, u);
}
for(int i = 1; i <= n; ++i)
lg[i] = lg[i - 1] + (1 << lg[i - 1] == i);
LCA_DFS(1, 0);

while(m--) {
int l, r, k;
cin >> l >> r >> k;
int lca = LCA(l, r);
// root[l] + root[r] - root[lca] - root[fa[lca][0]]前缀和思想
cout << vec[query(root[l], root[r], root[lca], root[fa[lca][0]], 1, n, k) - 1] << endl;
}
}

signed main() {
ios_base::sync_with_stdio(false);
//cin.tie(nullptr);
//cout.tie(nullptr);
##ifdef FZT_ACM_LOCAL
freopen("in.txt", "r", stdin);
freopen("out.txt", "w", stdout);
signed test_index_for_debug = 1;
char acm_local_for_debug = 0;
do {
if (acm_local_for_debug == '$') exit(0);
if (test_index_for_debug > 20)
throw runtime_error("Check the stdin!!!");
auto start_clock_for_debug = clock();
solve();
auto end_clock_for_debug = clock();
cout << "Test " << test_index_for_debug << " successful" << endl;
cerr << "Test " << test_index_for_debug++ << " Run Time: "
<< double(end_clock_for_debug - start_clock_for_debug) / CLOCKS_PER_SEC << "s" << endl;
cout << "--------------------------------------------------" << endl;
} while (cin >> acm_local_for_debug && cin.putback(acm_local_for_debug));
##else
solve();
##endif
return 0;
}

本文作者:jujimeizuo
本文地址https://blog.jujimeizuo.cn/2020/12/04/spoj-cot-count-on-a-tree-树上主席树-lca/
本博客所有文章除特别声明外,均采用 CC BY-SA 3.0 协议。转载请注明出处!