215. Kth Largest Element in an Array

Kesa...大约 7 分钟algorithm

215. 数组中的第K个最大元素open in new window

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

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

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

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

示例 2:

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

提示：

  • 1 <= k <= nums.length <= 10^5
  • -104 <= nums[i] <= 10^4

1. 快速选择 (快速排序分区)

快速选择使用的是快速排序的分区算法, 快速排序的分区算法有两种:

  • Lumoto Partition
  • Hoare Partition

关于基准的选择, 随机选取选取基准可以尽量避开极端情况.

1.1 Lomuto Partition

快速排序流程:

  1. 选取基准 Pivot
  2. 将小于于Pivot的数字放在其左边, 大于等于Pivot的放在右边, 形成左右分区; 左分区所有数字小于基准, 右分区所有数字大于等于基准; 对于左右分区, 继续 步骤1). 直到区间为空为止.
  3. 所有数字排序完毕

分区流程:

  1. 选取最右边的数作为基准
  2. 左分区索引i, 遍历数组, 将小于基准的数字和左分区数字交换; 直到结束为止;
  3. 将左分区边界右边的数字和基准交换, 返回当前基准位置

伪代码如下:

// Sorts a (portion of an) array, divides it into partitions, then sorts those
algorithm quicksort(A, lo, hi) is 
  // Ensure indices are in correct order
  if lo >= hi || lo < 0 then 
    return
    
  // Partition array and get the pivot index
  p := partition(A, lo, hi) 
      
  // Sort the two partitions
  quicksort(A, lo, p - 1) // Left side of pivot
  quicksort(A, p + 1, hi) // Right side of pivot

// Divides array into two partitions
algorithm partition(A, lo, hi) is 
  pivot := A[hi] // Choose the last element as the pivot

  // Temporary pivot index
  i := lo - 1

  for j := lo to hi - 1 do 
    // If the current element is less than or equal to the pivot
    if A[j] <= pivot then 
      // Move the temporary pivot index forward
      i := i + 1
      // Swap the current element with the element at the temporary pivot index
      swap A[i] with A[j]

  // Move the pivot element to the correct pivot position (between the smaller and larger elements)
  i := i + 1
  swap A[i] with A[hi]
  return i // the pivot index

由于 Lumoto Partition 返回的是当前基准所在位置, 那么一定能保证基准左右两边数字的大小关系.

  • 若基准位置正好为 n-k, 则表示其右边的k-1个数字比基准大, 那么基准就是第k大的数字.
  • 若基准位置小于 n-k, 表示第k大数字在其若右分区, 继续对右分区进行Partition, 寻找基准.
  • 若基准位置大于 n-k, 表示第k大数字在其若左分区, 继续对左分区进行Partition, 寻找基准.

每次只需要寻找一个分区即可, 无需等待数组排序完毕. 题解如下(Lumoto 分区法现在会超时, 因为加了新的测试用例):

func findKthLargest(nums []int, k int) int {
    n := len(nums)
    pi := lomutoPartition(nums, 0, n-1)
    t := n - k
    for pi != t {
        if pi < t {
            pi = lomutoPartition(nums, pi+1, n-1)
        } else {
            pi = lomutoPartition(nums, 0, pi-1)
        }
    }

    return nums[n-k]
}

func lomutoPartition(nums []int, left, right int) int {
    // random pivot
    pi := rand.Intn(right-left+1) + left
    p := nums[pi]
    // move to end
    swap(nums, pi, right)

    // partition
    i, j := left-1, left
    for j < right {
        if nums[j] < p {
            i++
            swap(nums, i, j)
        }
        j++
    }
    
    i++
    swap(nums, i, right) // move pivot to right position
    return i
}

func swap(nums []int, i, j int) {
    nums[i], nums[j] = nums[j], nums[i]
}

1.2 Hoare Partition

快速排序流程:

  1. 选取基准 Pivot
  2. 将小于于Pivot的数字放在其左边, 大于等于Pivot的放在右边, 形成左右分区; 左分区所有数字小于基准, 右分区所有数字大于等于基准; 对于左右分区, 继续 步骤1). 直到区间为空为止.
  3. 所有数字排序完毕

分区流程:

  1. 选择中间数字为基准

  2. 使用双指针i, j; 分别从数组两端向中间移动, 每次都移动一步:

    • 若nums[i] < pivot, 则i继续向右移动, 跳过这些数字
    • 若nums[j] > pivot, 则j继续向左移动, 跳过这些数字

    当i<j, 则交换. 直到i >= j

  3. 返回j.

算法伪代码:

// Sorts a (portion of an) array, divides it into partitions, then sorts those
algorithm quicksort(A, lo, hi) is 
  if lo >= 0 && hi >= 0 && lo < hi then
    p := partition(A, lo, hi) 
    quicksort(A, lo, p) // Note: the pivot is now included
    quicksort(A, p + 1, hi) 

// Divides array into two partitions
algorithm partition(A, lo, hi) is 
  // Pivot value
  pivot := A[ floor((hi - lo)/2) + lo ] // The value in the middle of the array

  // Left index
  i := lo - 1 

  // Right index
  j := hi + 1

  loop forever 
    // Move the left index to the right at least once and while the element at
    // the left index is less than the pivot
    do i := i + 1 while A[i] < pivot
    
    // Move the right index to the left at least once and while the element at
    // the right index is greater than the pivot
    do j := j - 1 while A[j] > pivot

    // If the indices crossed, return
    if i >= j then return j
    
    // Swap the elements at the left and right indices
    swap A[i] with A[j]

由于 Hoare Partition 返回的值j不一定是基准位置, 所以要使用递归的方式计算, 直到区间长度为1, 此时nums[n-k]为第k大数.

  1. 若j < n-k , 那么第k大在右分区, 递归调用 partition(nums, j+1, right) ,
  2. 若j >= n-k, 那么第k大在左分区, 递归调用 partition(nums, left, j) , 注意: 和 lomuto 不同, 左分区包含j自身; 所以左分区的条件一定是 j >= n-k
func findKthLargest(nums []int, k int) int {
    n := len(nums)
    return hoarePartition(nums, 0, n-1, n-k)
}  

func hoarePartition(nums []int, left, right, t int) int {
    if left == right {
        return nums[t]
    }

    pi := rand.Intn(right-left+1) + left
    p := nums[pi]

    i, j := left-1, right+1
    for i < j {
        for i++; nums[i] < p; {
            i++
        }
        for j--; nums[j] > p; {
            j--
        }
        if i < j {
            nums[i], nums[j] = nums[j], nums[i]
        }
    }

    if j < t {
        return hoarePartition(nums, j+1, right, t)
    } else {
        return hoarePartition(nums, left, j, t)
    }
}

2. 堆排序

因为需要求第K大的数, 那么结果是K个最大数字中最小的那个. 流程如下:

  1. 构建最小堆
  2. 遍历数组存入数字:
    • 若长度小于k, 直接存储
    • 若长度等于k, 比较堆顶元素, 若大于堆顶, 则入堆.
  3. 堆顶元素即位第K大数.

Golang 标准库可以使用heap包实现堆的操作, 结构体需要实现heap.Interface接口:

type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

sort.Interface:

type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a transitive ordering:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a transitive ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

题解:

func findKthLargest(nums []int, k int) int {
    // minIntHeap 
    h := &MinIntHeap{}
    
    for _, n := range nums {
        if h.Len() < k {
            heap.Push(h, n)
        } else if h.Len() == k {
            if (*h)[0] < n {
                heap.Pop(h)
                heap.Push(h, n)
            }
        }
    }

    return (*h)[0]
}

type MinIntHeap []int

func (h *MinIntHeap) Len() int {
    return len(*h)
}

func (h *MinIntHeap) Less(i, j int) bool {
    return (*h)[i] < (*h)[j]
}

func (h *MinIntHeap) Swap(i, j int) {
    (*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func (h *MinIntHeap) Push(x any) {
    *h = append(*h, x.(int))
}

func (h *MinIntHeap) Pop() any {
    oldLen := len(*h)
    val := (*h)[oldLen-1]
    *h = (*h)[: oldLen-1]
    return val
}

Reference

  1. https://en.wikipedia.org/wiki/Quicksort#Hoare_partition_schemeopen in new window
  2. https://pkg.go.dev/container/heap@go1.21.1#Interfaceopen in new window
上次编辑于:
0  字
评论
  • 按正序
  • 按倒序
  • 按热度
Powered by Waline v2.15.2