Java实现无需临时数组的快速选择算法查找第k小整数

核心思路 #

我们采用**原地三向切分（荷兰国旗分区）**逻辑替换原有的临时数组分区方案：直接在原数组的指定区间内完成元素分区，将区间内元素划分为小于基准值、等于基准值、大于基准值三个连续区域，递归时仅传递区间下标和k值，不需要创建任何临时数组，大幅降低空间开销。

具体实现代码 #

import java.util.Random;
import java.util.Scanner;

public class Main {
    private static Random rand = new Random();
    private static Scanner keyboard = new Scanner(System.in);

    // 对外暴露的调用入口，自动传入整个数组的左右边界
    public static int selectKthSmallest(int k, int[] S) {
        // k从1开始计数，数组区间默认是[0, S.length-1]
        return quickSelect(S, 0, S.length - 1, k);
    }

    // 核心递归方法，操作数组的[left, right]区间，查找该区间第k小的元素
    private static int quickSelect(int[] S, int left, int right, int k) {
        if (left == right) {
            return S[left];
        }

        // 随机选择基准值，避免最坏时间复杂度
        int pivotIndex = left + rand.nextInt(right - left + 1);
        int pivot = S[pivotIndex];

        // 三向切分指针初始化
        int lt = left; // 小于基准区的下一个插入位置，[left, lt-1]都是小于pivot的元素
        int gt = right; // 大于基准区的前一个插入位置，[gt+1, right]都是大于pivot的元素
        int i = left; // 当前遍历指针

        while (i <= gt) {
            if (S[i] < pivot) {
                swap(S, lt, i);
                lt++;
                i++;
            } else if (S[i] > pivot) {
                swap(S, i, gt);
                gt--;
            } else {
                i++;
            }
        }

        // 计算两个边界区间的长度
        int lessCount = lt - left; // 小于基准的元素总数
        int equalCount = gt - lt + 1; // 等于基准的元素总数

        if (k <= lessCount) {
            // 第k小在小于基准的区间，递归左半区
            return quickSelect(S, left, lt - 1, k);
        } else if (k <= lessCount + equalCount) {
            // 第k小在等于基准的区间，直接返回基准值
            return pivot;
        } else {
            // 第k小在大于基准的区间，k值减去前两个区间的总长度后递归右半区
            return quickSelect(S, gt + 1, right, k - lessCount - equalCount);
        }
    }

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

    public static void printArray(int[] S) {
        for (int i = 0; i < S.length; i++) {
            System.out.print(S[i] + "\t");
            if (i % 10 == 9)
                System.out.println();
        }
    }

    public static void main(String[] args) {
        System.out.print("Enter the size of an array: ");
        int size = keyboard.nextInt();
        while (size < 1) {
            System.out.println("Size of the array should be greater than 0.");
            System.out.print("Enter the size of an array: ");
            size = keyboard.nextInt();
        }

        System.out.print("Enter the value of k: ");
        int k = keyboard.nextInt();
        while (k < 1 || k > size) {
            System.out.println("Value of k should be in the range 1-" + size + ".");
            System.out.print("Enter the value of k: ");
            k = keyboard.nextInt();
        }

        int[] S = new int[size];
        for (int i = 0; i < size; i++) {
            S[i] = 100 + rand.nextInt(900);
        }

        System.out.println("\nRandom values generated in the array:");
        printArray(S);
        System.out.println();
        System.out.println(k + "th smallest value in the array: " + selectKthSmallest(k, S));
    }
}

优化说明 #

• 全程没有创建任何临时数组，仅通过交换原数组元素完成分区，额外空间开销仅为递归栈的O(logn)（平均情况）
• 三向切分的逻辑和原算法的S1/S2/S3分区逻辑完全等价，最终计算结果和原有实现一致
• 保留了随机选择基准值的逻辑，避免最坏情况O(n²)的时间复杂度

内容的提问来源于stack exchange，提问作者Kim Gordon