215.数组中的第K个最大元素(Medium)

题目描述

在未排序的数组中找到第 k 个最大的元素

样例

Input: [3,2,1,5,6,4] 和 k = 2
Output: 5

题解

方法一： 堆排

本题最适合的 O$(nlog_2n)$ 级的排序方法为堆排，小根堆满足题目所需的要求

方法二：快速选择

快速选择和快速排序的基本思路一致，不过只需要找到第 k 大的枢（pivot）即可，不需要对其左右再进行排序。

Python代码

方法一：堆排--Python自带的堆结构，默认小根堆

import heapq
from random import shuffle

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        shuffle(nums)  # 随机排序
        nums = [-1 * num for num in nums] # 小根堆求最大值
        heapq.heapify(nums) # 建立堆结构
        for _ in range(k - 1):
            heapq.heappop(nums)
        return -1 * heapq.heappop(nums)
        

方法二：快速排序

# 快速选择写法1
def quickSelection(nums, l, r):
    pivot = nums[l]
    i, j = l + 1, r 
    while True: # do..while Python实现
        while i < r and nums[i] <= pivot:
            i += 1
        while j > l and nums[j] >= pivot:
            j -= 1
        if i >= j:
            break
        nums[i], nums[j] = nums[j], nums[i]
    nums[l], nums[j] = nums[j], nums[l] # !!!!!! 必须是j
    return j

# 快速选择写法2
def quickSelection(nums, l, r):
    pivot = nums[l]
    low, high = l, r
    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 
  

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        l, r = 0, len(nums) - 1
        target = len(nums) - k
        while l < r:
            mid = quickSelection(nums, l, r)
            if mid == target:
                return nums[mid]
            elif mid < target:
                l = mid + 1
            else:
                r = mid - 1
        return nums[l]

Go 示例

Go 大根堆实现

// 创建堆
func build_heap(nums []int, n int) {
    for i := n / 2; i >= 0; i-- {
        adjust_heap(nums, i, n)
    }
}

// 调整堆
func adjust_heap(nums []int, i, n int) {
    lchild := 2 * i + 1
    rchild := 2 * i + 2
    max := i 
    if i < n / 2 {
        if lchild < n && nums[lchild] > nums[max] { max = lchild}
        if rchild < n && nums[rchild] > nums[max] { max = rchild}
        if max != i { // 如果确实调整了，就一直递归下去
            nums[max], nums[i] = nums[i], nums[max]
            adjust_heap(nums, max, n)
        }
    }
}

// 弹出堆的最大元素
func heap_pop(nums *[]int) int {
    pop_value := (*nums)[0]
    n := len(*nums)
    (*nums)[0] = (*nums)[n - 1]  // 最后一个元素和第一个进行交换 
    *nums = (*nums)[:n - 1]
    adjust_heap(*nums, 0, n - 1)
    return pop_value
}

func findKthLargest(nums []int, k int) int {
    n := len(nums)
    build_heap(nums, n)
    for i := 0; i < k - 1; i++ { // 弹出 k - 1 个最大的值
        heap_pop(&nums)
    }
    return nums[0]
}

Go 快排实现

func findKthLargest(nums []int, k int) int {
    n := len(nums)
    l, r := 0, n - 1
    for l < r {
        mid := quickSelect(nums, l, r)
        if mid == n - k {
            return nums[mid]
        } else if mid < n - k {
            l = mid + 1
        } else {
            r  = mid - 1
        }
    }
    return nums[l]
}

func quickSelect(nums []int, left, right int) int {
    pivot := nums[left]
    i, j := left + 1, right 
    for ;; {
        for left < j && nums[j] >= pivot {
            j--
        }
        
        for i < right && nums[i] <= pivot {
            i++
        }
        if i >= j {
            break
        }
        nums[i], nums[j] = nums[j], nums[i]
    }
    nums[left], nums[j] = nums[j], nums[left] // 一定是 j
    return j
}