数据结构 |堆 |问题 12

问题概述

本问题涉及到在一个大小为N的数组中查找第K大的元素。数组元素可以是任意类型，例如int, double等。

解决方案

方法一：使用堆排序

1. 创建一个大小为K的最小堆，堆中保存最大的K个元素。初始时堆中元素为数组前K个元素。
2. 遍历数组中剩余的元素，如果当前元素比堆中最小元素大，则将最小元素替换为当前元素，然后重新构建堆。
3. 遍历完数组后，堆顶元素即为第K大的元素。

代码示例：

public static <T extends Comparable<? super T>> T findKthLargest(T[] arr, int k) {
    PriorityQueue<T> pq = new PriorityQueue<>(k);
    for (T t : arr) {
        pq.offer(t);
        if (pq.size() > k) {
            pq.poll();
        }
    }
    return pq.peek();
}

方法二：使用快速选择算法

快速选择算法类似于快速排序算法，但是只需要对数组的一部分进行排序。

1. 选定一个基准元素pivot，把数组中所有小于pivot的元素都放在它的左边，所有大于pivot的元素都放在右边。
2. 判断pivot的位置是否为第K大的位置，如果是，则返回pivot。如果不是，则根据pivot的位置调整搜索范围：如果pivot的位置小于K，说明第K大的元素在右边，否则在左边。
3. 重复步骤1，直到找到第K大的元素。

代码示例：

public static <T extends Comparable<? super T>> T findKthLargest(T[] arr, int k) {
    int left = 0;
    int right = arr.length - 1;
    while (left <= right) {
        int pivotIndex = partition(arr, left, right);
        if (pivotIndex == k - 1) {
            return arr[pivotIndex];
        } else if (pivotIndex < k - 1) {
            left = pivotIndex + 1;
        } else {
            right = pivotIndex - 1;
        }
    }
    return null;
}

private static <T extends Comparable<? super T>> int partition(T[] arr, int leftBound, int rightBound) {
    T pivot = arr[leftBound];
    int left = leftBound + 1;
    int right = rightBound;
    while (left <= right) {
        while (left <= right && arr[left].compareTo(pivot) >= 0) {
            left++;
        }
        while (left <= right && arr[right].compareTo(pivot) < 0) {
            right--;
        }
        if (left < right) {
            swap(arr, left, right);
        }
    }
    swap(arr, leftBound, right);
    return right;
}

private static <T> void swap(T[] arr, int i, int j) {
    T temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

时间复杂度

方法一的时间复杂度为O(N log K)，其中N为数组长度，K为需要查找的第K大元素的位置。

方法二的时间复杂度为O(N)，其中N为数组长度。但是最坏情况下的时间复杂度为O(N^2)，因为在分割数组时，如果每次选定的基准元素都是最小值或最大值，会导致递归树的高度达到N，时间复杂度退化为O(N^2)。

空间复杂度

方法一的空间复杂度为O(K)，因为需要用大小为K的堆来存储最大的K个元素。

方法二的空间复杂度为O(1)，因为是原地排序，没有使用额外的空间。

适用场景

方法一适用于需要找到最大的K个元素的情况，例如找到前K大的员工薪水，或者找到前K大的股票涨幅。如果只需要查找第K大的元素，使用方法二更加高效。

方法二不仅适用于找到第K大的元素，还适用于找到第K小的元素。经过优化的快速选择算法可以在平均情况下实现O(N)的时间复杂度，并且不需要额外的空间。但是最坏情况下的时间复杂度较差，需要注意。