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

狙击美佐

于 2020-12-04 16:42:41 发布

阅读量200

点赞数 4

分类专栏：线段树主席树平衡树

本文链接：https://blog.csdn.net/fztsilly/article/details/110657097

版权

线段树主席树平衡树专栏收录该内容

10 篇文章 0 订阅

订阅专栏

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

前置技能

$可持久化线段树、 L C A$

题意

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

思路

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

$那么怎么能求第 k 值呢？$

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

$与树上两个点有联系的点，那就是\red{LCA--最近公共祖先}。$

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

$然后就是查询部分，根据下图：$

在这里插入图片描述

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

$在递归过程中，比较 h j t [h j t [q l] . l] . v a l + h j t [h j t [q r] . l] . v a l - h j t [h j t [l c a] . l] . v a l - h j t [h j t [f a l c a] . l] . v a l 与 k 的大小来判断要进入左子树还是右子树。$

Code

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