描述

347. 前 K 个高频元素

Difficulty: 中等

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

示例 1:

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

示例 2:

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

提示：

  • 你可以假设给定的 k 总是合理的，且 $1\leq k\leq m$, m 为数组中不相同的元素的个数。
  • 你的算法的时间复杂度 必须 优于 $O (nlogn)$ , n 是数组的大小。
  • 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的。
  • 你可以按任意顺序返回答案。

题解

思路

遍历数组，记录每个数字出现的次数，保存在哈希表中，这样的时间复杂度为 $O (n)$。

我们只需要找出这个哈希表中值最大的 k 个。

在这里我们可以用堆来排序。建立一个小顶堆，遍历出现次数数组。

  • 堆元素个数小于 k，入堆
  • 堆元素等于 k，检查堆顶与当前出现次数的大小，堆顶更大，则至少有 k 个数字出现次数比当前值大，舍弃当前值，否则弹出堆顶，插入当前值到堆。

遍历完成后，堆中元素纪委出现次数数组前 k 大的值。

遍历数组，用哈希表的记录时间为 $O (n)$，遍历出现次数数组，建立大小为 k 的堆，每次堆操作为时间复杂度为 $O (logk)$，n 次操作为 $O (nlogk)$，总时间复杂度为 $O (n)+O (nlogk)=O (n)$

代码

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
func topKFrequent(nums []int, k int) []int {
    occurrences := map[int]int{}
    for _, num := range nums {
        occurrences [num]++
    }
    h := &IHeap{}
    heap.Init (h)
    for key, value := range occurrences {
        heap.Push (h, [2]int{key, value})
        if h.Len () > k {
            heap.Pop (h)
        }
    }
    ret := make([]int, k)
    for i := 0; i < k; i++ {
        ret [k - i - 1] = heap.Pop (h).([2]int)[0]
    }
    return ret
}

type IHeap [][2]int

func (h IHeap) Len () int           { return len(h) }
func (h IHeap) Less (i, j int) bool { return h [i][1] < h [j][1] }
func (h IHeap) Swap (i, j int)      { h [i], h [j] = h [j], h [i] }

func (h *IHeap) Push (x interface{}) {
    *h = append(*h, x.([2]int))
}