编辑代码

import java.util.Random;

public class QuickSelect {
    public static void main(String[] args) {
        int[] arr = {3, 1, 4, 2, 5, 6};
        int k = 3; // 查找第 3 大元素
        int result = quickSelect(arr, k);
        System.out.println("第 " + k + " 大元素是: " + result);
    }

    public static int quickSelect(int[] arr, int k) {
        return quickSelect(arr, 0, arr.length - 1, k);
    }

    private static int quickSelect(int[] arr, int left, int right, int k) {
        if (left == right) {
            return arr[left];
        }

        int pivotIndex = partition(arr, left, right);
        int length = pivotIndex - left + 1; // 计算右侧部分的元素数量

        if (k == length) {
            return arr[pivotIndex];
        } else if (k < length) {
            return quickSelect(arr, left, pivotIndex - 1, k);
        } else {
            return quickSelect(arr, pivotIndex + 1, right, k - length);
        }
    }

    private static int partition(int[] arr, int left, int right) {
        // 随机选择一个基准元素
        Random rand = new Random();
        int randomIndex = rand.nextInt(right - left + 1) + left;
        int temp = arr[randomIndex];
        arr[randomIndex] = arr[right];
        arr[right] = temp;

        int pivot = arr[right];
        int i = left - 1;

        for (int j = left; j < right; j++) {
            if (arr[j] <= pivot) {
                i++;
                int temp2 = arr[i];
                arr[i] = arr[j];
                arr[j] = temp2;
            }
        }

        int temp3 = arr[i + 1];
        arr[i + 1] = arr[right];
        arr[right] = temp3;
        return i + 1;
    }
}