记录一个基于Java的利用快排切分来实现快排TopK问题的代码模板

使用快排切分实现快排和TopK问题的解题模板

import java.util.Arrays;

public class TestDemo {
    public static void main(String[] args) {
        int[] arr1 = {48, 12, 6 ,8, 11};
        int[] arr2 = {25, 14, 45 ,8, 10};
        int k = 2;
        System.out.print("Array_1: ");
        printArr(arr1);
        System.out.print("GetTopK: ");
        // 注意是k-1
        int[] topK = getTopK(arr1, 0, arr1.length - 1, k - 1);
        printArr(topK);
        System.out.print("Array_2: ");
        printArr(arr2);
        quickSort(arr2, 0, arr2.length - 1);
        System.out.print("Sorted:  ");
        printArr(arr2);
    }

    /**
     * 快速排序
     * @param arr 数组
     * @param l 左边界
     * @param r 右边界
     */
    private static void quickSort(int[] arr, int l, int r) {
        if (l >= r) return;
        int index = partition(arr, l, r);
        quickSort(arr, l, index - 1);
        quickSort(arr, index + 1, r);
    }

    /**
     * 快排切分寻找TopK
     * @param arr 数组
     * @param l 左边界
     * @param r 右边界
     * @param k TopK
     * @return TopK数组
     */
    private static int[] getTopK(int[] arr, int l, int r, int k) {
        // 每快排切分1次，找到排序后下标为j的元素，如果j恰好等于k就返回j以及j左边所有的数
        int j = partition(arr, l, r);
        if (j == k) return Arrays.copyOf(arr, j + 1);
        // 否则根据下标j与k的大小关系来决定继续切分左段还是右段
        return j > k ? getTopK(arr, l, j - 1, k) : getTopK(arr, j + 1, r, k);
    }

    /**
     * 快排切分
     * 返回下标j，使得比arr[j]小的数都在j的左边，比arr[j]大的数都在j的右边
     * @param arr 数组
     * @param l 左边界
     * @param r 右边界
     * @return 一次切分后的基准值下标j
     */
    private static int partition(int[] arr, int l, int r) {
        int i = l, j = r + 1, pivot = arr[l];
        while (true) {
            while (++i <= r && arr[i] < pivot);
            while (--j >= l && arr[j] > pivot);
            if (i >= j) break;
            swap(arr, i, j);
        }
        swap(arr, l, j);
        return j;
    }

    /**
     * 数组下标对应值交换函数
     * @param arr 数组
     * @param i 下标i
     * @param j 下标j
     */
    private static void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    /**
     * 打印数组
     * @param arr 数组
     */
    private static void printArr(int[] arr) {
        for (int i : arr) {
            System.out.print(i + " ");
        }
        System.out.println();
    }
}

输出结果

Array_1: 48 12 6 8 11
GetTopK: 6 8
Array_2: 25 14 45 8 10
Sorted:  8 10 14 25 45