如何实现Top K算法？

Top K算法是一种快速查找数组中前K大元素的算法。常见的 Top K 算法有:

1. 基于快速排序的方法。使用快速排序将数组分割为K大元素所在的数组。时间复杂度O(nlogK)。
2. 基于堆的方法。维护一个大小为K的最小堆,遍历数组将元素插入堆中,并在插入后调整堆。时间复杂度O(nlogK)。
3. 基于快速选择的方法。使用快速选择算法找到数组第K大元素。时间复杂度O(n)。

我们这里实现第3种基于快速选择的Top K算法:

public int[] topK(int[] nums, int k) {
    int[] result = new int[k];
    quickselect(nums, 0, nums.length - 1, k, result);
    return result;
}

private void quickselect(int[] nums, int left, int right, int k, int[] result) {
    if (left == right) return;

    int pivotIndex = partition(nums, left, right);
    if (pivotIndex > k - 1) {
        quickselect(nums, left, pivotIndex - 1, k, result);
    } else if (pivotIndex < k - 1) {
        for (int i = left; i <= pivotIndex; i++) {
            result[i - left] = nums[i];
        }
        quickselect(nums, pivotIndex + 1, right, k - (pivotIndex - left + 1), result);
    } else {
        for (int i = left; i <= right; i++) {
            result[i - left] = nums[i];
        }
    }
}

private int partition(int[] nums, int left, int right) {
    // 选择最后一个元素为pivot
    int pivot = nums[right];
    int j = left;
    for (int i = left; i < right; i++) {
        if (nums[i] <= pivot) {
            swap(nums, i, j);
            j++;
        }
    }
    swap(nums, j, right);
    return j;
}

private void swap(int[] nums, int i, int j) {
    int temp = nums[i];
    nums[i] = nums[j];
    nums[j] = temp;
}

算法流程:

1. 递归终止:当left == right,返回。
2. 选择pivot并划分数组,返回pivot位置index。
3. 如果pivotIndex > k – 1,递归查找左部分top k。
4. 如果pivotIndex < k – 1, result数组前pivotIndex – left + 1个元素为top k,递归查找右部分top k – (pivotIndex – left + 1)。
5. 否则找到top k,返回。
6. 不断递归划分与查找,直到找到top k。

时间复杂度O(n),空间复杂度O(k)。

所以,Top K算法的关键是:

1. 基于快速选择算法,使用pivot划分数组。
2. 递归终止条件:left == right。
3. 根据pivotIndex判断top k在左部分、右部分或找到,相应递归或返回。
4. 通过不断递归划分与查找,找到top k。

与快速选择算法相比,找到top k后直接返回result,不继续递归。