[spoj COT - Count on a tree]樹上第K小

[spoj COT - Count on a tree]樹上第K小

分類：Data Structure Presidental Tree template

1. 題目鏈接

[spoj COT - Count on a tree]

2. 題意描述

N個節點的樹，樹上每個節點有一個點權。M次詢問，每次詢問一條鏈上的第k 小數。
數據範圍：(N,M<=100000)
Time limit: 0.129s-0.516s
Memory limit: 1536MB

3. 解題思路

本來以爲跟COT2一樣可以用莫隊在樹上跑。第k小的維護可以用pbds，線段樹之類的東西維護。這樣的複雜度是O(nn√∗log2n) 。本題時間卡得比較緊，而且上述算法的複雜度偏高。
正解應該是主席樹。看到給出的內存限制這麼大應該聯想到了...
在序列上，求第k 小，步驟是先將點進行離散化，用主席樹維護的是一個離散化之後點的前綴和。然後，查詢的時候，對區間二分。
將情形推廣到樹上，同樣的也是先將點進行離散化，然後從根節點，dfs向下依次將節點加入到主席樹中。
那麼詢問(u→v) 的路徑上的點的時候，就是sum[u]+sum[v]−2∗sum[lca(u,v)]+lca(u,v) .

4. 實現代碼

#include <bits/stdc++.h>

using namespace std;

#define FIN(x) freopen(#x".txt", "r", stdin)

const int MAXN = 110000;
const int DEEP = 20;

int n, m;

int w[MAXN], f[MAXN], fsz;

struct Edge {
    int v, next;
} edge[MAXN << 1];
int head[MAXN], tot, dep[MAXN], fa[MAXN][DEEP];
void init_edge() {
    tot = 0;
    memset(head, -1, sizeof(head));
}
inline void add_edge(int u, int v) {
    edge[tot] = Edge{v, head[u]};
    head[u] = tot ++;
}

int lca(int u, int v) {
    while(dep[u] != dep[v]) {
        if(dep[u] < dep[v]) swap(u, v);
        int d = dep[u] - dep[v];
        for(int i = 0; i < DEEP; ++i) {
            if(d >> i & 1) u = fa[u][i];
        }
    }
    if(u == v) return u;
    for(int i = DEEP - 1; i >= 0; --i) {
        if(fa[u][i] != fa[v][i]) {
            u = fa[u][i];
            v = fa[v][i];
        }
    }
    return fa[u][0];
}

#define lch(u)  nd[(u)].ch[0]
#define rch(u)  nd[(u)].ch[1]
struct TNode {
    int sum, ch[2];
    inline void reset() { sum = ch[0] = ch[1] = 0; }
} nd[MAXN * 20];
int ndIdx, root[MAXN];

void build(int l, int r, int& rt) {
    rt = ++ ndIdx;
    nd[rt].reset();
    if(l == r) return;
    int md = (l + r) >> 1;
    build(l, md, lch(rt));
    build(md + 1, r, rch(rt));
}

void update(int p, int l, int r, int& rt, int prt) {
    rt = ++ ndIdx;
    nd[rt] = nd[prt];
    ++ nd[rt].sum;
    if(l == r) return;
    int md = (l + r) >> 1;
    if(p <= md) update(p, l, md, lch(rt), lch(prt));
    else update(p, md + 1, r, rch(rt), rch(prt));
}

int kth(int k, int l, int r, int u, int v, int anc, int ancp) {
    if(l == r) return l;
    int md = (l + r) >> 1;
    int sum = nd[lch(u)].sum + nd[lch(v)].sum - nd[lch(anc)].sum * 2 + (l <= ancp && ancp <= md);
    if(sum >= k) return kth(k, l, md, lch(u), lch(v), lch(anc), ancp);
    else return kth(k - sum, md + 1, r, rch(u), rch(v), rch(anc), ancp);
}

int kth(int k, int u, int v) {
    int anc = lca(u, v), ancp = w[anc];
    int le = 1, ri = fsz, md, sum;
    u = root[u], v = root[v], anc = root[anc];
    while(le < ri) {
        md = (le + ri) >> 1;
        sum = nd[lch(u)].sum + nd[lch(v)].sum - nd[lch(anc)].sum * 2 + (le <= ancp && ancp <= md);
        if(sum >= k) {
            ri = md;
            u = lch(u), v = lch(v), anc = lch(anc);
        } else {
            k -= sum;
            le = md + 1;
            u = rch(u), v = rch(v), anc = rch(anc);
        }
    }
    return le;
}

void dfs(int u, int pre, int deep) {
    int v;
    dep[u] = deep; fa[u][0] = pre;
    update(w[u], 1, fsz, root[u], root[pre]);
    for(int i = head[u]; ~i; i = edge[i].next) {
        v = edge[i].v;
        if(v == pre) continue;
        dfs(v, u, deep + 1);
    }
}

int main() {
#ifdef ___LOCAL_WONZY___
    FIN(input);
#endif // ___LOCAL_WONZY___
    int u, v, k;
    scanf("%d %d", &n, &m);
    fsz = 0;
    for(int i = 1; i <= n; ++i) {
        scanf("%d", &w[i]);
        f[++ fsz] = w[i];
    }
    sort(f + 1, f + fsz + 1);
    fsz = unique(f + 1, f + fsz + 1) - f - 1;
    for(int i = 1; i <= n; ++i) w[i] = lower_bound(f + 1, f + fsz + 1, w[i]) - f;

    init_edge();
    for(int i = 2; i <= n; ++i) {
        scanf("%d %d", &u, &v);
        add_edge(u, v);
        add_edge(v, u);
    }
    ndIdx = 0;
    build(1, fsz, root[0]);
    dfs(1, 0, 0);
    for(int i = 1; i < DEEP; i++) {
        for(int j = 1; j <= n; j++) {
            fa[j][i] = fa[fa[j][i - 1]][i - 1];
        }
    }
    for(int i = 1; i <= m; ++i) {
        scanf("%d %d %d", &u, &v, &k);
        int pos = kth(k, u, v);
//        int anc = lca(u, v);
//        int pos = kth(k, 1, fsz, root[u], root[v], root[anc], w[anc]);
        printf("%d\n", f[pos]);
    }
#ifdef ___LOCAL_WONZY___
    cout << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC * 1000 << "ms." << endl;
#endif // ___LOCAL_WONZY___
    return 0;
}