leetcode 347. 前 K 个高频元素 golang实现

描述

给定一个非空的整数数组，返回其中出现频率前k高的元素。

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]
示例 2:

输入: nums = [1], k = 1
输出: [1]
说明：

你可以假设给定的k总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
你的算法的时间复杂度必须优于 O(n log n) ,n是数组的大小。

思路

1. 遍历数组 map记录数字出现的次数
2. 利用长度k的最小堆排序，root节点始终存放的为出现次数最小的数字
3. 遍历map 与堆root比较，如果大于root出现的次数，则将新的数字替换掉root 再次排序
4. 遍历完成后堆里保存的即为结果

实现

type element struct{
	Num   int   // 元素数字
	Times int   // 数字出现的次数
}

func topKFrequent(nums []int, k int) []int {
    // 遍历数组 记录数字出现的次数
	numTimeMap := make(map[int]int, 0)
	for _, num := range nums{
		numTimeMap[num]++
	}

	heap := make([]element, 0)              // 堆数组
	ret := make([]int, 0)                   // 返回结果
	sortFlag := false                       // 判断堆数组长度是否达到k 的标记位
	for num, times := range numTimeMap{
		e := element{Num:num, Times:times}
		if len(heap) < k{                   // 堆数组没有达到k 则继续append
			heap = append(heap, e)
			continue
		}else if !sortFlag{                 // 堆数组达到k并且还为重建堆数组 则重建堆数组
			topKBuildHeap(heap)
			sortFlag = true
		}

		if e.Times > heap[0].Times{         // 如果数字出现的次数大于堆数组root节点数字的出现次数 则将root节点替换 并重新排序
			heap[0] = e
			topKHeapSortBody(heap, k, 0)

		}
	}

	for _, e := range heap{
		ret = append(ret, e.Num)
	}

	return ret
}

// 重新堆数组
func topKBuildHeap(elementList []element){
	size := len(elementList)

	for i := size; i >= 0; i--{
		topKHeapSortBody(elementList, size, i)
	}

	return
}

// 堆排序
func topKHeapSortBody(elementList []element, size, curRoot int){
	left := 2*curRoot+1
	right := 2*curRoot+2
	smallest := curRoot


	if left < size && elementList[smallest].Times > elementList[left].Times {
		smallest = left
	}

	if right < size && elementList[smallest].Times > elementList[right].Times {
		smallest = right
	}

	if smallest != curRoot{
		elementList[smallest], elementList[curRoot] = elementList[curRoot], elementList[smallest]
		topKHeapSortBody(elementList, size, smallest)
	}

	return
}