poj 2104 K-th Number

Problem

poj.org/problem?id=2104
vjudge.net/problem/POJ-2104

Reference

主席树介绍
树状结构之主席树

Meaning

给一个序列 $a_{n}$，有 m 次询问，每次询问区间 [ L , R ] 中第 k 大的数是多少。

Analysis

主席树模板题。注意空间要开足，数组越界报了很多次 WA 而不是 RE…

Notes

主席树逻辑上是多棵线段树。这道题序列长度为 n，将数据离散化之后，$a_{i}$ 的范围就会变成 [ 1 , n ]，第 i 棵线段树中的结点（范围是[ l , r ]）记录的是 在 $a_{1}$ ~ $a_{i}$ 中处于 [ l , r ] 这个范围的元素个数。
其它介绍见参考博客。

Code

#include <cstdio>
#include <algorithm>
using namespace std;
const int N = 100000;

struct node
{
    int v;
    int l, r; // 左右儿子
} tree[20*N];

int a[N+1]; // 原序列
int b[N+1]; // 离散化数组
int root[N+1], sz; // 第i棵线段树的树根的位置

int build(int l, int r)
{
    int rt = sz++;
    tree[rt].v = tree[rt].l = tree[rt].r = 0;
    if(l != r)
    {
        int m = l + r >> 1;
        tree[rt].l = build(l, m);
        tree[rt].r = build(m+1, r);
    }
    return rt;
}

int add_link(int v, int l, int r, int pre)
{
    int rt = sz++;
    tree[rt] = tree[pre];
    ++tree[rt].v;
    if(l != r)
    {
        int m = l + r >> 1;
        if(v > m)
            tree[rt].r = add_link(v, m+1, r, tree[pre].r);
        else
            tree[rt].l = add_link(v, l, m, tree[pre].l);
    }
    return rt;
}

int query(int ql, int qr, int k, int l, int r)
{
    if(l == r)
        return l;
    int m = l + r >> 1;
    int left = tree[tree[qr].l].v - tree[tree[ql].l].v;
    if(left < k)
        return query(tree[ql].r, tree[qr].r, k-left, m+1, r);
    else
        return query(tree[ql].l, tree[qr].l, k, l, m);
}

int main()
{
    int n, m;
    scanf("%d%d", &n, &m);
    for(int i=1; i<=n; ++i)
    {
        scanf("%d", a+i);
        b[i] = a[i];
    }
    sort(b + 1, b + n + 1);
    int top = unique(b + 1, b + n + 1) - b - 1;
    sz = 0;
    root[0] = build(1, top);
    for(int i=1, v; i<=n; ++i)
    {
        v = lower_bound(b + 1, b + n + 1, a[i]) - b;
        root[i] = add_link(v, 1, top, root[i-1]);
    }
    for(int l, r, k; m--; )
    {
        scanf("%d%d%d", &l, &r, &k);
        int x = query(root[l-1], root[r], k, 1, top);
        printf("%d\n", b[x]);
    }
    return 0;
}

一开始建的那棵空树，只需要用一个结点就可以了，让它的左右儿子也指向 0 号位置，那指来指去就都是那个 0 号空结点，于是一棵树可以缩成一个点，建树函数build()可以改成：

void build()
{
    root[0] = tree[0].v = tree[0].l = tree[0].r = sz++;
}