TopK 算法 - 小顶堆

概述

题目：Leetcode 703 : 取出数据流中第K大的元素
描述：设计一个找到数据流中第 k 大元素的类（class）。注意是排序后的第 k 大元素，不是第 k 个不同的元素。请实现 KthLargest 类：
- KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
- int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。

思路

使用数组 nums []int 存储数据流，模拟堆结构，堆中只存储 k 个元素，且堆顶元素nums[0]为最小值。
每次插入数据，如果数据流中的元素个数小于 k，直接插入堆中，如果数据流中的元素个数大于等于 k，判断插入的元素是否大于堆顶元素，如果大于堆顶元素，则将堆顶元素替换为插入的元素，然后重新调整堆，保持堆的有序性。

插入元素时，执行上浮算法将该元素与父级元素比较和交换，直到父级元素小于等于该元素。
替换堆首元素时，执行下沉算法，将堆首元素与子级元素比较和交换，直到子级元素大于该元素。

获取第 k 大元素，直接返回 nums[0] 即可。

实现

初始化数据结构

定义结构体类型 minHeap 并创建一个堆结构

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


// 堆结构
type minHeap struct {
 nums []int  // 存储数据
 k    int  // 堆中元素个数，非实际堆中元素个数
}

// createMinHeap 创建一个堆
func createMinHeap(nums []int, k int) *minHeap {
 // 初始化
 heap := &minHeap{nums: []int{}, k: k}

 for i := 0; i < len(nums); i++ {
  heap.add(nums[i])
 }
 return heap
}

在minHeap上定义add方法，用于插入元素，并保持堆的有序性。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


func (mh *minHeap) add(num int) int {
 if mh.k > len(mh.nums) { // 堆中还有位置
  mh.nums = append(mh.nums, num) // 添加在堆尾，然后开始上浮
  mh.swim(len(mh.nums) - 1)      // 上浮
 } else { // 堆中没有位置
  if num > mh.nums[0] { // num 大于堆顶的数字，即大于堆中的最小值，将最小值替换为num，然后下沉
   mh.nums[0] = num
   mh.sink(0) // 下沉
  }
 }

 return mh.nums[0]
}

定义下沉算法sink

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


// sink 下沉
func (mh *minHeap) sink(idx int) {
 // idx 开始下沉的位置
 for 2*idx+1 < len(mh.nums) { // 左子节点不溢出
  // 当前节点与左右两个节点中较小的一个交换
  child := 2*idx + 1
  if child+1 < len(mh.nums) && mh.nums[child] > mh.nums[child+1] {
   child++ // 取右子节点
  }

  if mh.nums[child] >= mh.nums[idx] {
   // 子节点大于等于当前值，下沉结束
   break
  }

  // 交换
  mh.nums[idx], mh.nums[child] = mh.nums[child], mh.nums[idx]
  // 将下标更新为当前交换的子节点
  idx = child
 }
}

定义上浮算法swim

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


// swim 上浮
func (mh *minHeap) swim(idx int) {
 // idx 开始上浮的位置
 for idx > 0 { // idx 上浮至0(根节点)时停止
  parent := (idx - 1) >> 1 // 父节点
  if mh.nums[parent] <= mh.nums[idx] {
   // 父节点小于等于当前值，上浮结束
   break
  }

  // 交换
  mh.nums[parent], mh.nums[idx] = mh.nums[idx], mh.nums[parent]
  idx = parent
 }
}

– 完整代码如下

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88



package main

// 小顶堆
type minHeap struct {
 nums []int
 k    int
}

func createMinHeap(nums []int, k int) *minHeap {
 // 初始化
 heap := &minHeap{nums: []int{}, k: k}

 for i := 0; i < len(nums); i++ {
  heap.add(nums[i])
 }
 return heap
}

// add 添加元素
func (mh *minHeap) add(num int) int {
 if mh.k > len(mh.nums) { // 堆中还有位置
  mh.nums = append(mh.nums, num) // 添加在堆尾，然后开始上浮
  mh.swim(len(mh.nums) - 1)      // 上浮
 } else { // 堆中没有位置
  if num > mh.nums[0] { // num 大于堆顶的数字，即大于堆中的最小值，将最小值替换为num，然后下沉
   mh.nums[0] = num
   mh.sink(0) // 下沉
  }
 }

 return mh.nums[0]
}

// sink 下沉
func (mh *minHeap) sink(idx int) {
 // idx 开始下沉的位置
 for 2*idx+1 < len(mh.nums) { // 左子节点不溢出
  // 当前节点与左右两个节点中较小的一个交换
  child := 2*idx + 1
  if child+1 < len(mh.nums) && mh.nums[child] > mh.nums[child+1] {
   child++ // 取右子节点
  }

  if mh.nums[child] >= mh.nums[idx] {
   // 子节点大于等于当前值，下沉结束
   break
  }

  // 交换
  mh.nums[idx], mh.nums[child] = mh.nums[child], mh.nums[idx]
  // 将下标更新为当前交换的子节点
  idx = child
 }
}

// swim 上浮
func (mh *minHeap) swim(idx int) {
 // idx 开始上浮的位置
 for idx > 0 { // idx 上浮至0时停止
  parent := (idx - 1) >> 1 // 父节点
  if mh.nums[parent] <= mh.nums[idx] {
   // 父节点小于等于当前值，上浮结束
   break
  }

  // 交换
  mh.nums[parent], mh.nums[idx] = mh.nums[idx], mh.nums[parent]
  idx = parent
 }
}

type KthLargest struct {
 heap *minHeap
}

// Constructor 初始化
func Constructor(k int, nums []int) KthLargest {
 return KthLargest{heap: createMinHeap(nums, k)}
}

func (this *KthLargest) Add(val int) int {
 return this.heap.add(val)
}

func main() {
  k := Constructor(3, []int{4, 5, 8, 2})
}