347.M 前 K 个高频元素

Problem: 347. 前 K 个高频元素

思路

思路1：大根堆

先统计每个数字的个数，然后提取出唯一的数字，然后构造大根堆，再不断取大根堆的根节点元素，以及缩小堆的大小，这样取出来的结果是按频率降序排列的。

这种思路的问题是，如果去重后的数字很多的话，构建堆的成本会很高。

思路2：小根堆

先统计每个数字的个数，然后构造小根堆。

当堆中的元素个数小于k时，追加新元素到堆中，并且向上调整新追加的元素；否则，如果新元素大于根节点，则替换根节点，并且向下调整堆。

Code

代码1：

func topKFrequent(nums []int, k int) []int {
	countMap := make(map[int]int)
	for _, num := range nums {
		countMap[num]++
	}

	var uniqNums []int
	for num := range countMap {
		uniqNums = append(uniqNums, num)
	}

	for i := len(uniqNums)/2 - 1; i >= 0; i-- {
		maxHeapify(uniqNums, countMap, i, len(uniqNums))
	}

	var ans []int
	for i := 0; i < k; i++ {
		ans = append(ans, uniqNums[0])
		uniqNums[0], uniqNums[len(uniqNums)-i-1] = uniqNums[len(uniqNums)-i-1], uniqNums[0]
		maxHeapify(uniqNums, countMap, 0, len(uniqNums)-i-1)
	}
	return ans
}

func maxHeapify(nums []int, countMap map[int]int, root int, heapSize int) {
	l, r, largest := 2*root+1, 2*root+2, root
	if l < heapSize && countMap[nums[l]] > countMap[nums[largest]] {
		largest = l
	}

	if r < heapSize && countMap[nums[r]] > countMap[nums[largest]] {
		largest = r
	}

	if largest != root {
		nums[largest], nums[root] = nums[root], nums[largest]
		maxHeapify(nums, countMap, largest, heapSize)
	}
}

代码2：

func topKFrequent(nums []int, k int) []int {
	freq := make(map[int]int)
	for _, num := range nums {
		freq[num]++
	}

	var heap []int
	for num := range freq {
		if len(heap) < k {
			heap = append(heap, num)
			siftUp(heap, freq, len(heap)-1)
		} else {
			if freq[num] > freq[heap[0]] {
				heap[0] = num
				siftDown(heap, freq, 0, len(heap))
			}
		}
	}

	return heap
}

func siftUp(nums []int, freq map[int]int, index int) {
	parent := (index - 1) / 2
	if parent >= 0 && freq[nums[index]] < freq[nums[parent]] {
		nums[parent], nums[index] = nums[index], nums[parent]
		siftUp(nums, freq, parent)
	}
}

func siftDown(nums []int, freq map[int]int, root int, heapSize int) {
	l, r, smallest := 2*root+1, 2*root+2, root
	if l < heapSize && freq[nums[l]] < freq[nums[smallest]] {
		smallest = l
	}

	if r < heapSize && freq[nums[r]] < freq[nums[smallest]] {
		smallest = r
	}

	if smallest != root {
		nums[smallest], nums[root] = nums[root], nums[smallest]
		siftDown(nums, freq, smallest, heapSize)
	}
}