347. 前 K 个高频元素

120 阅读1分钟

题目:
给你一个整数数组 nums 和一个整数 k ,请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。
解法:
堆的性质

func topKFrequent(nums []int, k int) []int {
	
	numsCount := make(map[int]int)
	for i := range nums {
		numsCount[nums[i]] ++
	}
	heap := NewHeap(k)
	for num, count := range numsCount {
		heap.insert([2]int{num, count})
	}

	ans := make([]int, 0)

	for i := 1; i <= heap.N; i ++ {
		ans = append(ans, heap.Array[i][0])
	}
	return ans
}

// root节点最小
// Array的第0位不存数据,索引k的父节点位置为k/2,子节点位置为2k,2k+1
type Heap struct{
	Array [][2]int
	N int
	Capacity int
}

func NewHeap(n int) *Heap{
	return &Heap{
		Array: make([][2]int, 1),
		Capacity: n,
	}
}

func (h *Heap)insert(v [2]int){
	if h.N < h.Capacity {
		h.Array = append(h.Array, v)
		h.N ++
		h.swim(h.N)	
	} else {
		if v[1] > h.Array[1][1] {
			h.Array[1] = v
			h.sink(1)
		}
	}
	
}


// // 删除根节点的值
// func (h *Heap)deleteRoot() int{
// 	ret := h.Array[1]
// 	h.swap(1, h.N)
// 	h.Array = h.Array[:h.N]
// 	h.N --
// 	h.sink(1)
// 	return ret
// }

// 节点i的父节点比它小时,执行swim,节点i上浮,
func (h *Heap)sink(i int){
	current := i
	child := 2 * current
	for child <= h.N {
		// 右子节点一定比左子节点大,交换右子节点和根节点一定满足堆的条件
		if child + 1 <= h.N && h.less(child, child + 1) {
			child = child + 1
		}
		if !h.less(current, child){
			break
		}
		h.swap(current, child)
		current = child
		child = 2 * current
	}
}

// 节点i的子节点比它大时,执行sink,节点i下沉
func (h *Heap)swim(i int){
	current := i
	parent := current / 2
	for parent > 0 {
		if !h.less(parent, current) {
			break
		}
		h.swap(current, parent)

		current = parent
		parent = current / 2
	}
}

func (h *Heap)less(i int, j int) bool {
	return h.Array[i][1] > h.Array[j][1]
}

func (h *Heap)swap(i int, j int)  {
	h.Array[i], h.Array[j] = h.Array[j], h.Array[i]
}