每日算法 — 使用java实现快速选择:基于QuickSelect的第K大元素查找与三数取中优化

在刷题和实际工程中,「查找第K大(或第K小)元素」是一个高频需求:从海量用户中找出积分排名前100的玩家、在日志系统中定位第K大的错误级别、实时统计直播间的礼物榜TopK……如果先排序再取第K个,时间复杂度是O(n log n)。而快速选择(QuickSelect)算法基于快速排序的分区思想,能够在平均O(n)时间内直接定位目标元素,是处理此类问题的利器。本文将用Java完整实现QuickSelect算法,并加入三数取中(Median-of-Three)优化来避免最坏情况。

算法核心思想

QuickSelect由Tony Hoare发明,与快速排序(QuickSort)师出同门。其核心思路是:

  1. 分区(Partition):像快排一样选定一个基准值(pivot),将数组划分为「小于pivot」和「大于pivot」的两个区域,并确定pivot的最终位置索引p。
  2. 比较与递归:将p与目标索引k比较:
  3. 若p == k,pivot就是第K小元素,直接返回。
  4. 若p > k,目标元素在左半区,对左半区递归。
  5. 若p < k,目标元素在右半区,对右半区递归。

由于每次只需处理一半数据,平均时间复杂度降至O(n)。但在极端数据(如已排序数组)下,若总是选第一个或最后一个元素作为pivot,会退化为O(n²)。因此本文采用三数取中法:从数组头、中、尾三个位置取中位数作为pivot,大幅提升稳定性。

Java完整实现

以下是完整的QuickSelect实现,包含三数取中优化、单指针分区策略,以及查找第K大/第K小的封装方法。

import java.util.Random;

/**
 * QuickSelect算法实现:在O(n)平均时间内查找无序数组中的第K大或第K小元素。
 * 采用三数取中(Median-of-Three)策略选择pivot,避免在有序数据上退化。
 */
public class QuickSelect {

    private static final Random RANDOM = new Random();

    /**
     * 查找数组中第K大的元素(1-based,k=1表示最大元素)。
     *
     * @param nums 输入数组
     * @param k    第几大(1 <= k <= nums.length)
     * @return 第K大的元素值
     */
    public static int findKthLargest(int[] nums, int k) {
        if (nums == null || k < 1 || k > nums.length) {
            throw new IllegalArgumentException("Invalid input: k must be between 1 and array length.");
        }
        // 第K大等价于第(nums.length - k)小(0-based索引)
        return quickSelect(nums, 0, nums.length - 1, nums.length - k);
    }

    /**
     * 查找数组中第K小的元素(1-based,k=1表示最小元素)。
     *
     * @param nums 输入数组
     * @param k    第几小(1 <= k <= nums.length)
     * @return 第K小的元素值
     */
    public static int findKthSmallest(int[] nums, int k) {
        if (nums == null || k < 1 || k > nums.length) {
            throw new IllegalArgumentException("Invalid input: k must be between 1 and array length.");
        }
        // 第K小等价于索引(k - 1)处的元素
        return quickSelect(nums, 0, nums.length - 1, k - 1);
    }

    /**
     * QuickSelect核心递归方法。
     * 在nums[left..right]范围内查找最终位于索引targetIdx处的元素。
     *
     * @param nums       输入数组(会被修改)
     * @param left       当前搜索区间左边界
     * @param right      当前搜索区间右边界
     * @param targetIdx  目标索引(0-based)
     * @return targetIdx处的元素值
     */
    private static int quickSelect(int[] nums, int left, int right, int targetIdx) {
        while (left <= right) {
            // 使用三数取中法选择pivot,并将其放到right-1位置
            int pivot = medianOfThree(nums, left, right);

            // 分区:将数组分为 < pivot | == pivot | > pivot 三部分
            // 返回pivot的最终位置范围 [pivotLeft, pivotRight]
            int[] pivotRange = partition3Way(nums, left, right, pivot);

            if (targetIdx >= pivotRange[0] && targetIdx <= pivotRange[1]) {
                // 目标索引落在pivot的相等区域内,直接返回
                return nums[targetIdx];
            } else if (targetIdx < pivotRange[0]) {
                // 目标在左半区(小于pivot的部分)
                right = pivotRange[0] - 1;
            } else {
                // 目标在右半区(大于pivot的部分)
                left = pivotRange[1] + 1;
            }
        }
        // 理论上不会执行到这里,因为targetIdx始终在有效范围内
        throw new IllegalStateException("Target index out of bounds during selection.");
    }

    /**
     * 三数取中:从left、mid、right三个位置中选择中位数作为pivot。
     * 同时将三个数排序,使得 nums[left] <= nums[mid] <= nums[right]。
     * 返回选定的pivot值(即nums[mid])。
     *
     * @param nums  数组
     * @param left  左索引
     * @param right 右索引
     * @return pivot值
     */
    private static int medianOfThree(int[] nums, int left, int right) {
        int mid = left + (right - left) / 2;

        // 将三个元素排序:确保 nums[left] <= nums[mid] <= nums[right]
        if (nums[left] > nums[mid]) {
            swap(nums, left, mid);
        }
        if (nums[left] > nums[right]) {
            swap(nums, left, right);
        }
        if (nums[mid] > nums[right]) {
            swap(nums, mid, right);
        }

        // 将中位数放到 right - 1 的位置,作为pivot的隔离区
        // 这样 left 位置一定 <= pivot,right 位置一定 >= pivot
        swap(nums, mid, right - 1);
        return nums[right - 1];
    }

    /**
     * 三路分区(3-Way Partition):将数组分为 < pivot、== pivot、> pivot 三部分。
     * 这对于含有大量重复元素的数据集非常高效。
     *
     * @param nums   数组
     * @param left   左边界
     * @param right  右边界
     * @param pivot  基准值
     * @return int[2],其中 [0] 是等于pivot区域的左边界,[1] 是右边界
     */
    private static int[] partition3Way(int[] nums, int left, int right, int pivot) {
        int lt = left;      // lt 指向 < pivot 区域的下一个位置
        int gt = right;     // gt 指向 > pivot 区域的前一个位置
        int i = left;       // 当前扫描指针

        while (i <= gt) {
            if (nums[i] < pivot) {
                // 当前元素小于pivot,放到lt区域
                swap(nums, i, lt);
                lt++;
                i++;
            } else if (nums[i] > pivot) {
                // 当前元素大于pivot,放到gt区域
                swap(nums, i, gt);
                gt--; // i不递增,因为交换过来的元素需要重新判断
            } else {
                // 当前元素等于pivot,留在中间区域
                i++;
            }
        }

        // 返回等于pivot的区域范围 [lt, gt]
        return new int[]{lt, gt};
    }

    /**
     * 交换数组中两个位置的元素。
     */
    private static void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }

    // ==================== 测试与演示 ====================

    public static void main(String[] args) {
        // 测试用例1:普通数组,查找第3大元素
        int[] arr1 = {3, 2, 1, 5, 6, 4};
        System.out.println("Array: [3, 2, 1, 5, 6, 4]");
        System.out.println("3rd largest: " + findKthLargest(arr1, 3)); // 预期: 4
        System.out.println("2nd smallest: " + findKthSmallest(arr1.clone(), 2)); // 预期: 2
        System.out.println();

        // 测试用例2:含有大量重复元素
        int[] arr2 = {3, 2, 3, 1, 2, 4, 5, 5, 6};
        System.out.println("Array: [3, 2, 3, 1, 2, 4, 5, 5, 6]");
        System.out.println("4th largest: " + findKthLargest(arr2, 4)); // 预期: 4
        System.out.println();

        // 测试用例3:已排序数组(验证三数取中防退化效果)
        int[] arr3 = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
        System.out.println("Sorted array: [1..10]");
        System.out.println("1st largest: " + findKthLargest(arr3, 1)); // 预期: 10
        System.out.println("10th largest (smallest): " + findKthLargest(arr3.clone(), 10)); // 预期: 1
        System.out.println();

        // 测试用例4:逆序数组
        int[] arr4 = {10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
        System.out.println("Reverse sorted array: [10..1]");
        System.out.println("5th largest: " + findKthLargest(arr4, 5)); // 预期: 6
        System.out.println();

        // 测试用例5:大规模随机数据性能测试
        int n = 1_000_000;
        int[] largeArr = new int[n];
        for (int i = 0; i < n; i++) {
            largeArr[i] = RANDOM.nextInt(n);
        }
        long start = System.currentTimeMillis();
        int result = findKthLargest(largeArr, 500_000);
        long elapsed = System.currentTimeMillis() - start;
        System.out.println("1M random elements, find 500000th largest: " + result);
        System.out.println("Time elapsed: " + elapsed + "ms");
    }
}

三路分区的优势

上述实现采用了三路分区(3-Way Partition),这是处理含大量重复元素数据集的关键优化。传统单指针分区在面对重复元素时效率不高,而三路分区将数组划分为三个区域:

  • < pivot 的区域
  • == pivot 的区域
  • > pivot 的区域

当目标索引恰好落在== pivot的区域内时,可以立即返回结果,无需继续递归。这使得算法在极端重复数据下依然保持线性时间复杂度。

复杂度分析

指标 平均情况 最坏情况(未优化) 最坏情况(三数取中)
时间复杂度 O(n) O(n²) 接近O(n),极难触发退化
空间复杂度 O(1) O(n)(递归栈) O(log n)(递归栈)

三数取中法将最坏情况的发生概率降至极低。对于已排序或逆序数组,传统QuickSelect会退化,但三数取中总能选出中间值作为pivot,保证每次大致二分。

实际应用场景

  1. 实时排行榜:游戏服务器需要在千万玩家中实时获取前100名,可用QuickSelect先定位第100大的分数线,再遍历一遍收集所有大于等于该分数的玩家。
  2. 数据流异常检测:在日志系统中,用QuickSelect维护最近一段时间内第95百分位的响应时间,作为服务健康度的阈值。
  3. 推荐系统冷启动:为新用户快速从海量商品中筛选出评分中位数的商品作为初始推荐池。
  4. A/B测试显著性检验:在多组实验数据中,用QuickSelect快速获取中位数和四分位数,评估实验效果。

总结

QuickSelect是处理「第K大/小」问题的经典算法,平均O(n)的时间复杂度和O(1)的额外空间使其在工程实践中极具价值。通过三数取中法选择pivot,我们有效避免了最坏情况;通过三路分区,我们进一步优化了重复元素场景下的性能。掌握这一算法,能让你在面对TopK、中位数查找等问题时,写出既高效又优雅的代码。