学习笔记-划分树 + 查询给定范围内第K大的数

刚开始不是很理解划分树存在的意义，因为查询第K大的数并不难，只要把数组排序，然后取第K大的数不就行了吗? 为什么要搞个划分树这个东西出来呢? 后面在网上查了一下之后，发现划分树还是有其存在的意义的.考虑下面的例子:
给你n个数的原序列，有m次询问，每次询问给出l、r、k，求原序列l到r之间第k 大的数。n范围10万，m范围5千.
像这种题目，就很适合用划分树来解决问题.上面的题有两个关键点，第一是'每次查询都是原序列的l到r之间查询'.第二个关键点是，可能会有m次查询，也就是说查询的次数可能会很多.
如果用快排的思想来处理的话，因为排序会改变原数组的顺序，因此每次查询，都需要创建一个新的数组来存放排序的后的数据以保证不改变原数组的顺序.因此快排算法的空间复杂度会很高。在时间复杂度上，快排也并没有优势.
如果用划分树的话就不一样了，只需要一次建树的过程，之后的查询都是在已建好的树中进行查询.因此时间复杂度是m * logn + 建树的时间.当m的次数很大的情况下，建树的时间就可以忽略不计了.
当然如果m的次数很少，那么也可以考虑使用快排的思路来解决问题.

代码实现

/**
 * 划分树 + 查询.
 */
public class PartitionTreeSearch {
    // 第一个维度代表层数.
    private int tree[][];
    // 已排序的数组.
    private int sorted[];
    // 记录被放入左子树的个数.
    // 比如num[deep][i]，表示第deep层，i之前有多少数字被放入了deep+1层的左子树中.
    // num[0][2] = 1，表示第0层的下标2之之前有一个数字被放入了第0层的左子树中。
    private int num[][];
    private int source[];

    public PartitionTreeSearch(int[] source) {
        build(source);
    }

    public void build(int[] source) {
        init(source);
        build(0, source.length - 1, 0);
    }

    /**
     * 查询下标范围在rl到rr之间的第K大的数字.
     *
     * @param rl
     * @param rr
     * @param k
     * @return
     */
    public int query(int rl, int rr, int k) {
        return query(0, 0, source.length - 1, rl, rr, k);
    }

    public int query(int deep, int L, int R, int rl, int rr, int K) {
        int M = (L + R) >> 1;
        if (rl == rr) return tree[deep][rl];

        int lToLeft;
        int rToLeft;
        if (rl == L) {
            lToLeft = 0;
            rToLeft = num[deep][rr];
        } else {
            lToLeft = num[deep][rl - 1];
            rToLeft = num[deep][rr] - lToLeft;// rl和rr之间有多少个去了左子树.
        }
        // 如果去左子树的个数大于等于K的话，说明第K大的数肯定在左子树，因此直接进入左子树.
        if (rToLeft >= K) {
            int nrl = L + lToLeft;// 新的rl边界，公式为L + rl之前的进入左子树的数字的个数.
            int nrr = L + lToLeft + rToLeft - 1;// 新的rr边界，公式为L + rr之前进入左子树的数字的个数
            return query(deep + 1, L, M - 1, nrl, nrr, K);
        } else {
            // 括号表示rl之前有多少个数进入了右子树.
            int nrl = M + (rl - L - lToLeft) + 1;
            // 括号表示rr之前有多少个数进入了右子树，包括rr本身.
            int nrr = M + (rr - L - rToLeft - lToLeft + 1);
            // 表示在右子树查第K - rToLeft的数.
            return query(deep + 1, M + 1, R, nrl, nrr, K - rToLeft);
        }
    }

    private void init(int[] source) {
        this.source = source;
        tree = new int[20][source.length];
        num = new int[20][source.length];
        sorted = new int[source.length];
        for (int i = 0; i < source.length; i++) {
            tree[0][i] = source[i];
            sorted[i] = source[i];
        }

        Arrays.sort(sorted);
    }

    private void build(int L, int R, int deep) {
        if (L >= R) return;
        int M = (L + R) >> 1; // 取中间下标.
        int isSame = M - L + 1; // 保存和mid相等，但是被放入左子树的数的个数.
        int mid = sorted[M]; // 注意这里取的是排序过后的数组的中位数.

        for (int i = L; i <= R; i++) {
            // 左子树本来可以放M - L + 1个数，
            // 所有小于mid的数都需要被放入左子树，当isSame减去这些数之后，isSame表示的意思就是左子树需要存放多少与mid相等的数字.
            if (tree[deep][i] < mid) {
                isSame--;
            }
        }

        int l = L;
        int r = M + 1;
        for (int i = L; i <= R; i++) {
            int v = tree[deep][i];
            // 初始化.
            if (i == L) {
                num[deep][i] = 0;
            } else {
                num[deep][i] = num[deep][i - 1];
            }

            if (v < mid) {
                tree[deep + 1][l++] = v;
                num[deep][i]++;
            } else if (v == mid && isSame > 0) {
                tree[deep + 1][l++] = v;
                isSame--;//左子树可存放与mid相等的数的个数减1
                num[deep][i]++;
            } else {
                tree[deep + 1][r++] = v;
            }
        }

        build(L, M, deep + 1);
        build(M + 1, R, deep + 1);
    }
}