SPOJ COT (主席树 LCA)

题目链接:点击打开链接

题意:求树上u-v之间路径上的第k小的节点值.

和数列的第k小类似不过写起来烦一点.在数列上求区间第k小是两个前缀减一减,

树上就是u-root,v-root前缀加一加减去两倍LCA-root,新建线段树都是在父亲节

点的基础上更新的.

#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <map>
#include <set>
#include <vector>
#include <queue>
using namespace std;
#define maxn 211111
#define maxm maxn*40

int a[maxn], n, m, q;
vector <int> num;
map <int , int> gg;
struct E {
    int v, next;
}edge[maxn];
int head[maxn], cnt, Hash[maxn];

//主席树
struct node {
    int l, r, num;
}tree[maxm];
int tot, T[maxn];

int build_tree (int l, int r) {
    int root = tot++;
    tree[root].num = 0;
    if (l == r)
        return root;
    int mid = (l+r)>>1;
    tree[root].l = build_tree (l, mid);
    tree[root].r = build_tree (mid+1, r);
    return root;
}

int update (int root, int pos, int val) {
    int new_root = tot++;
    int tmp = new_root;
    tree[new_root].num = tree[root].num+val;
    int l = 1, r = m;
    while (l < r) {
        int mid = (l+r)>>1;
        if (pos <= mid) {
            tree[new_root].r = tree[root].r;
            tree[new_root].l = tot++;
            new_root = tree[new_root].l;
            root = tree[root].l;
            r = mid;
        }
        else {
            tree[new_root].l = tree[root].l;
            tree[new_root].r = tot++;
            new_root = tree[new_root].r;
            root = tree[root].r;
            l = mid+1;
        }
        tree[new_root].num = tree[root].num + val;
    }
    return tmp;
}

void dfs (int u, int fa) {
    T[u] = update (T[fa], a[u], 1);
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].v;
        if (v == fa) continue;
        dfs (v, u);
    }
}

void build () {
    tot = 0;
    T[n+1] = build_tree (1, m);
    dfs (1, 0);
}

int query (int l_root, int r_root, int root, int k) {//u,v,lca的主席树节点
    int pos = a[root];
    root = T[root];
    int l = 1, r = m;
    while (l < r) {
        int mid = (l+r)>>1;
        int tmp = tree[tree[l_root].l].num + tree[tree[r_root].l].num -
        tree[tree[root].l].num*2 + (pos >= l && pos <= mid);
        if (tmp >= k) {
            l_root = tree[l_root].l;
            r_root = tree[r_root].l;
            root = tree[root].l;
            r = mid;
        }
        else {
            k -= tmp;
            l_root = tree[l_root].r;
            r_root = tree[r_root].r;
            root = tree[root].r;
            l = mid+1;
        }
    }
    return l;
}

//LCA
int fa[maxn][22], deep[maxn];//节点i第2^j个祖先 深度
int DEG;
bool vis[maxn];

void bfs (int root) {
    DEG = 20;
    queue <int> q;
    while (!q.empty ()) q.pop ();
    deep[root] = 0;
    fa[root][0] = root;
    q.push (root);
    memset (vis, 0, sizeof vis); vis[root] = 1;
    while (!q.empty ()) {
        int u = q.front (); q.pop ();
        for (int i = 1; i < DEG; i++) {
            fa[u][i] = fa[fa[u][i-1]][i-1];
        }
        for (int i = head[u]; i != -1; i = edge[i].next) {
            int v = edge[i].v;
            if (vis[v]) continue;
            deep[v] = deep[u]+1;
            fa[v][0] = u;
            q.push (v);
            vis[v] = 1;
        }
    }
}

int LCA (int u, int v) {
    if (deep[u] > deep[v])
        swap (u, v);
    int hu = deep[u], hv = deep[v];
    int tu = u, tv = v;
    for (int det = hv-hu, i = 0; det; det >>= 1, i++) if (det&1) {
        tv = fa[tv][i];
    }
    if (tu == tv)
        return tu;
    for (int i = DEG-1; i >= 0; i--) {
        if (fa[tu][i] == fa[tv][i])
            continue;
        tu = fa[tu][i];
        tv = fa[tv][i];
    }
    return fa[tu][0];
}

void add_edge (int u, int v) {
    edge[cnt].v = v, edge[cnt].next = head[u], head[u] = cnt++;
}

void init () {//初始化
    gg.clear ();
    num.clear ();
    memset (head, -1, sizeof head);
    cnt = 0;
}

void lisan () {//离散化
    int cnt = 0;
    sort (num.begin (), num.end ());
    for (int i = 0; i < n; i++) {
        if (!i || num[i] != num[i-1]) {
            gg[num[i]] = ++cnt;
            Hash[cnt] = num[i];
        }
    }
    for (int i = 1; i <= n; i++) a[i] = gg[a[i]];
    m = cnt;
}

int main () {
    while (scanf ("%d%d", &n, &q) == 2) {
        init ();
        for (int i = 1; i <= n; i++) {
            scanf ("%d", &a[i]);
            num.push_back (a[i]);
        }
        lisan ();
        for (int i = 1; i < n; i++) {
            int u, v;
            scanf ("%d%d", &u, &v);
            add_edge (u, v);
            add_edge (v, u);
        }
        bfs (1);
        build ();
        for (int i = 1; i <= q; i++) {
            int u, v, k;
            scanf ("%d%d%d", &u, &v, &k);
            int fa = LCA (u, v);
            printf ("%d\n", Hash[query (T[u], T[v], fa, k)]);
        }
    }
    return 0;
}