HDU 2665 Kth Number：劃分樹（主席樹）模板題詳解

HDU 2665 Kth Number：劃分樹（主席樹）模板題詳解

這是一道經典的區間第 K 小查詢問題，標準解法為劃分樹（Partition Tree），也稱主席樹。給定一個長度為 n 的數列，需回答 m 個查詢：區間 [l, r] 中第 k 小的數是多少。

核心思想

劃分樹類似於歸並排序的過程，將數值域遞歸二分，並在每層記錄左半部分有多少數字。構建時，將當前區間的數值按中位數分為左右兩組，保持相對順序，並記錄每個位置分到左邊的數字個數（toleft 數組）。查詢時，根據 toleft 判斷第 k 小落在左還是右子區間，遞歸下降直到葉節點。

關鍵數據結構

• tree[dep][i]：第 dep 層第 i 個位置的值
• sorted[i]：原數組排序後的結果，用於確定分割值
• toleft[dep][i]：第 dep 層前 i 個位置中分到左子樹的個數

算法流程

1. 構建： build(l, r, dep) 遞歸劃分。計算 mid、same（等於分割值且分左的配額），將元素分配到左右子樹，更新 toleft。
2. 查詢： query(L, R, l, r, dep, k) 根據 toleft 計算左子區間個數 cnt，若 k ≤ cnt 遞歸左子樹，否則遞歸右子樹並扣減 k。

複雜度

構建 O(n log n)，單次查詢 O(log n)，空間 O(n log n)。n ≤ 10⁵ 完全足夠。

參考代碼 (C++)

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int MAXN = 100010;
int tree[20][MAXN];
int sorted[MAXN];
int toleft[20][MAXN];

void build(int l, int r, int dep) {
    if (l == r) return;
    int mid = (l + r) >> 1;
    int same = mid - l + 1;
    for (int i = l; i <= r; i++)
        if (tree[dep][i] < sorted[mid]) same--;
    int lpos = l, rpos = mid + 1;
    for (int i = l; i <= r; i++) {
        if (tree[dep][i] < sorted[mid])
            tree[dep+1][lpos++] = tree[dep][i];
        else if (tree[dep][i] == sorted[mid] && same > 0) {
            tree[dep+1][lpos++] = tree[dep][i];
            same--;
        } else
            tree[dep+1][rpos++] = tree[dep][i];
        toleft[dep][i] = toleft[dep][l-1] + lpos - l;
    }
    build(l, mid, dep+1);
    build(mid+1, r, dep+1);
}

int query(int L, int R, int l, int r, int dep, int k) {
    if (l == r) return tree[dep][l];
    int mid = (L + R) >> 1;
    int cnt = toleft[dep][r] - toleft[dep][l-1];
    if (cnt >= k) {
        int newl = L + toleft[dep][l-1] - toleft[dep][L-1];
        int newr = newl + cnt - 1;
        return query(L, mid, newl, newr, dep+1, k);
    } else {
        int newr = r + toleft[dep][R] - toleft[dep][r];
        int newl = newr - (r - l - cnt);
        return query(mid+1, R, newl, newr, dep+1, k - cnt);
    }
}

int main() {
    int T, n, m;
    scanf("%d", &T);
    while (T--) {
        memset(tree, 0, sizeof(tree));
        scanf("%d%d", &n, &m);
        for (int i = 1; i <= n; i++) {
            scanf("%d", &tree[0][i]);
            sorted[i] = tree[0][i];
        }
        sort(sorted+1, sorted+n+1);
        build(1, n, 0);
        while (m--) {
            int l, r, k;
            scanf("%d%d%d", &l, &r, &k);
            printf("%d\n", query(1, n, l, r, 0, k));
        }
    }
    return 0;
}