【题解】HackerRank-BST maintenance

题目

https://www.hackerrank.com/challenges/bst-maintenance/problem

题意

按照给定顺序往二叉查找树中添加结点,回答每次添加结点后的所有结点对的距离和。

题解

  • 首先线性构建笛卡尔树。(当然也有别的构建方法,但更麻烦一些。)
  • 考虑插入一个结点 \(u\),新增加的距离和为 \(\sum_v dep[u]+dep[v]-2\times dep[lca(u, v)]\),lca 的部分用树链剖分 + 线段树/树状数组维护。

代码

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
#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 N = 25E4 + 100;
int val[N], a[N];
int G[N][2];
int n, rt;

void build() {
static int s[N], last;
int p = 0;
FOR (x, 1, n + 1) {
last = 0;
while (p && val[s[p - 1]] > val[x]) last = s[--p];
if (p) G[s[p - 1]][1] = x;
if (last) G[x][0] = last;
s[p++] = x;
}
rt = s[0];
}

int fa[N], dep[N], idx[N], out[N], ridx[N];
namespace hld {
int sz[N], son[N], top[N], clk;
void predfs(int u, int d) {
dep[u] = d; sz[u] = 1;
int& maxs = son[u] = -1;
for (int& v: G[u]) {
if (!v || v == fa[u]) continue;
fa[v] = u;
predfs(v, d + 1);
sz[u] += sz[v];
if (maxs == -1 || sz[v] > sz[maxs]) maxs = v;
}
}
void dfs(int u, int tp) {
top[u] = tp; idx[u] = ++clk; ridx[clk] = u;
if (son[u] != -1) dfs(son[u], tp);
for (int& v: G[u])
if (v && v != fa[u] && v != son[u]) dfs(v, v);
out[u] = clk;
}
template<typename T>
int go(int u, int v, T&& f = [](int, int) {}) {
int uu = top[u], vv = top[v];
while (uu != vv) {
if (dep[uu] < dep[vv]) { swap(uu, vv); swap(u, v); }
f(idx[uu], idx[u]);
u = fa[uu]; uu = top[u];
}
if (dep[u] < dep[v]) swap(u, v);
if (u != v) f(idx[v] + 1, idx[u]);
return v;
}
}

namespace bit {
LL c[N], cc[N];
inline LL lowbit(LL x) { return x & -x; }
void add(LL x, LL v) {
for (LL i = x; i <= n; i += lowbit(i)) {
c[i] += v; cc[i] += x * v;
}
}
void add(LL l, LL r, LL v) { add(l, v); add(r + 1, -v); }
LL sum(LL x) {
LL ret = 0;
for (LL i = x; i > 0; i -= lowbit(i))
ret += (x + 1) * c[i] - cc[i];
return ret;
}
LL sum(LL l, LL r) { return sum(r) - sum(l - 1); }
}

int main() {
cin >> n;
FOR (i, 1, n + 1) {
int t; scanf("%d", &t);
a[i] = t; val[t] = i;
}
build();
hld::predfs(rt, 0);
hld::dfs(rt, rt);
LL ans = 0, dsum = 0;
FOR (i, 1, n + 1) {
int u = a[i];
hld::go(u, rt, [&](int l, int r){ ans -= 2 * bit::sum(l, r); });
ans += 1LL * (i - 1) * dep[u] + dsum;
printf("%lld\n", ans);
hld::go(u, rt, [&](int l, int r){ bit::add(l, r, 1); });
dsum += dep[u];
}
}