[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(n\sqrt{n}*log_2{n})$。本题时间卡得比较紧，而且上述算法的复杂度偏高。
正解应该是主席树。看到给出的内存限制这么大应该联想到了...
在序列上，求第$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;
}