数组中的第K个最大元素

题目链接

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例:

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5

思路

解法一：

使用快排

  1. 寻找到第k个数停止递归，使得nums数组中index左边是前k个小的数，index右边是后面n-k个大的数
  2. 利用partion函数进行排序
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        low, high = 0, len(nums) -1
        while True:
            index = self.partion(nums, low, high)
            if index == k - 1:
                return nums[index]
            elif index < k - 1: # 在右边
                low = index + 1
            else: # 在左边
                high = index - 1

    def partion(self, nums, low, high):
        pivot = nums[low]
        while low < high:
            while low < high and nums[high] <= pivot: high -= 1 # 若右指针元素小于等于基准值，右指针左移
            nums[low] = nums[high] # 交换元素
            while low < high and nums[low] >= pivot: low += 1 # 若左指针元素大于等于基准值，左指针右移
            nums[high] = nums[low] # 交换元素
        nums[low] = pivot # 将基准元素放入最终位置
        return low

解法二

该题很容易想到使用快速排序的方法，这里介绍一种常见的方法：堆排序。

堆排序中主要有两种操作，一种是称为上浮 swim 操作，一种称为下沉 sink 操作。

  • 上浮操作：将指定节点与父节点进行比较，若满足特定条件，则交换指定节点与父节点；
  • 下沉操作：将指定节点与左右子节点进行比较，首先找到左右子节点中满足条件的节点，然后与指定节点进行比较，若满足条件，则交换元素；

思路如下：

  1. 首先利用上浮操作创建一个大根堆；
  2. 将最后一个元素与根节点进行交换，对根节点进行堆结构的调整 k - 1 次。
class Solution
{
public:
    int findKthLargest(vector<int> &nums, int k)
    {
        int heapSize = nums.size();
        // 建成大根堆
        for (int i = heapSize / 2; i >= 0; i--)
        {
            sink(nums, i, heapSize);
        }
        for (int i = nums.size() - 1; i >= nums.size() - k + 1; i--)
        {
            swap(nums[0], nums[i]);
            heapSize--;
            sink(nums, 0, heapSize);
        }

        return nums[0];
    }

private:
    // 上浮 从下到上调整堆
    void swim(vector<int> &heap, int i)
    {
        while (i > 0 && heap[i] < heap[(i - 1) / 2])
        { // 与父节点进行比较
            swap(heap[i], heap[(i - 1) / 2]);
            i = (i - 1) / 2;
        }
    }

    // 下沉 从下到上调整堆
    void sink(vector<int> &heap, int i, int N)
    {
        while (2 * i + 1 < N)
        {
            int j = 2 * i + 1;
            if (j + 1 < N && heap[j + 1] > heap[j])
            { // 从两个节点中找到最大的元素
                j++;
            }
            if (heap[i] > heap[j])
            {
                break;
            }
            swap(heap[i], heap[j]);
            i = j;
        }
    }
};

解法三

该题很容易想到使用快速排序的方法，这里介绍一种常见的方法：堆排序。

  1. 从大到小依次对每个非叶子节点进行堆结构的调整，构成大根堆；
  2. 交换根节点和最后的叶子节点，即删除 k - 1 次元素
class Solution
{
public:
    int findKthLargest(vector<int> &nums, int k)
    {
        int heapSize = nums.size();
        buildMaxHeap(nums, heapSize);
        for (int i = nums.size() - 1; i >= nums.size() - k + 1; --i) // 调整 k - 1 次
        {
            swap(nums[0], nums[i]);
            --heapSize;
            maxHeapify(nums, 0, heapSize);  //重新调整堆
        }
    }

    void maxHeapify(vector<int> &nums, int index, int heapSize) // 对节点进行一次调整
    {
        int left = 2 * index + 1, right = 2 * index + 2, largest = index;
        if (left < heapSize && nums[left] > nums[largest])
        {
            largest = left;
        }
        if (right < heapSize && nums[right] > nums[largest])
        {
            largest = right;
        }
        if (largest != index) // 若找到的最大元素不是 index, 则交换元素
        {
            swap(nums[index], nums[largest]);
            maxHeapify(nums, largest, heapSize); // 递归调整
        }
    }

    void buildMaxHeap(vector<int> &nums, int heapSize)
    {
        for (int i = heapSize / 2; i >= 0; i--)
        {
            maxHeapify(nums, i, heapSize); // 创建大根堆
        }
    }
};