「SPOJ 10628」COT - Count on a Tree

题目链接:SPOJ 10628

你有一棵 $n$ 个节点的树,节点从 $1$ 到 $n$ 编号。每个点都有一个权值 $a_i$。现在有 $m$ 个询问,每个询问形如:

  • u v k:求节点 $u,v$ 之间的路径上的第 $k$ 小权值。

数据范围:$1\le n,m\le 10^5$。


Solution

这题只不过是把区间第 $k$ 小放到了树上。我们还是利用主席树。由于节点 $u,v$ 之间的路径可以通过若干从根出发的路径相减而得,于是我们维护 $n$ 棵线段树,第 $i$ 个点的线段树由其父亲继承而来。每次求出 $\text{LCA}$ 后按照套路在线段树上二分即可。

时间复杂度:$\mathcal O((n+m)\log n)$。


Code

#include <cstdio>
#include <cstring>
#include <algorithm>

const int N = 1e5 + 5, M = 2e5 + 5, logN = 18;

int n, m, sz, tot, a[N], b[N], lnk[N], ter[M], nxt[M], f[N][logN], dep[N];

struct Chairman {
    static const int M = 2e6 + 5;
    int idx, rt[N], sum[M], ls[M], rs[M];
    void build(int &p, int l, int r) {
        p = ++idx, sum[p] = 0;
        if (l == r) {
            return;
        }
        int mid = (l + r) >> 1;
        build(ls[p], l, mid);
        build(rs[p], mid + 1, r);
    }
    void modify(int &p, int l, int r, int u, int x, int v) {
        p = ++idx, ls[p] = ls[u], rs[p] = rs[u], sum[p] = sum[u] + v;
        if (l == r) {
            return;
        }
        int mid = (l + r) >> 1;
        if (x <= mid) {
            modify(ls[p], l, mid, ls[u], x, v);
        } else {
            modify(rs[p], mid + 1, r, rs[u], x, v);
        }
    }
    int query(int l, int r, int L, int R, int A, int F, int k) {
        if (l == r) {
            return l;
        }
        int mid = (l + r) >> 1;
        int val = sum[ls[L]] + sum[ls[R]] - sum[ls[A]] - sum[ls[F]];
        if (k <= val) {
            return query(l, mid, ls[L], ls[R], ls[A], ls[F], k);
        } else {
            return query(mid + 1, r, rs[L], rs[R], rs[A], rs[F], k - val);
        }
    }
} cha;

void add(int u, int v) {
    ter[++tot] = v, nxt[tot] = lnk[u], lnk[u] = tot; 
}
int discretize() {
    for (int i = 1; i <= n; i++) b[i] = a[i];
    std::sort(b + 1, b + n + 1);
    int sz = std::unique(b + 1, b + n + 1) - (b + 1);
    for (int i = 1; i <= n; i++) {
        a[i] = std::lower_bound(b + 1, b + sz + 1, a[i]) - b;
    }
    return sz;
}
void dfs(int u, int p) {
    f[u][0] = p, dep[u] = dep[p] + 1;
    cha.modify(cha.rt[u], 1, sz, cha.rt[p], a[u], 1);
    for (int i = 1; (1 << i) <= dep[u]; i++) {
        f[u][i] = f[f[u][i - 1]][i - 1];
    }
    for (int i = lnk[u]; i; i = nxt[i]) {
        int v = ter[i];
        if (v == p) continue;
        dfs(v, u);
    }
}
int lca(int u, int v) {
    if (dep[u] < dep[v]) std::swap(u, v);
    int d = dep[u] - dep[v];
    for (int i = 17; ~i; i--) {
        if (d & (1 << i)) u = f[u][i];
    }
    if (u == v) return u;
    for (int i = 17; ~i; i--) {
        if (f[u][i] != f[v][i]) u = f[u][i], v = f[v][i];
    }
    return f[u][0];
}
int main() {
    scanf("%d%d", &n, &m);
    for (int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
    }
    sz = discretize();
    for (int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        add(u, v), add(v, u);
    }
    cha.build(cha.rt[0], 1, sz);
    dfs(1, 0);
    for (int i = 1; i <= m; i++) {
        int u, v, k;
        scanf("%d%d%d", &u, &v, &k);
        int x = lca(u, v);
        int ans = cha.query(1, sz, cha.rt[u], cha.rt[v], cha.rt[x], cha.rt[f[x][0]], k);
        printf("%d\n", b[ans]);
    }
    return 0;
}

1 条评论

  1. fffxk

    %%%DSY

发表评论