【题解】HackerRank-BST maintenance

Posted on In Disqus:

题目

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
107
#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];
    }
}