Top K Frequent Elements
Top K Frequent Elements
好久没刷题了,做了道之前做过的 medium 难度的题 Top K Frequent Elements
Given a non-empty array of integers, return the k most frequent elements.
Example 1:
Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]
Example 2:
Input: nums = [1], k = 1
Output: [1]
Note:
You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
Your algorithm's time complexity must be better than O(n log n), where n is the array's size.
思路
分成两步来解决:
- 先用 Map 统计下每个数出现的次数, 剩下的就是一个 TopN 问题;
- 遍历 Map 中的键值对, 遍历过程中维护一个容量为K的小顶堆.
小顶堆维护过程如下:
- 如果小顶堆中元素不足K个, 将当前元素添加到堆中;
- 如果小顶堆中有K个元素, 并且当前 key 的出现次数大于堆顶的 key 的出现次数, 就用当前 key 替换堆顶 key, 并维护小顶堆性质.
代码
这里用 kotlin
来实现:
// Runtime 248ms Memory 42.5MB
class Solution {
fun topKFrequent(nums: IntArray, k: Int): List<Int> {
val map = HashMap<Int, Int>()
nums.forEach {
map[it] = (map[it]?:0) + 1
}
if (k >= map.keys.size) {
return map.keys.toList()
}
val result = ArrayList<Int>(k)
map.forEach { (key, value) ->
if (result.size < k) {
// 堆中 key 不足 k 个, 将当前 key 放入堆中
result.add(key)
// 维护小顶堆
var i = result.lastIndex
while (i > 0) {
val tmp = (i - 1) / 2
if (map[result[tmp]]!! <= value) {
break
} else {
result[i] = result[tmp]
result[tmp] = key
i = tmp
}
}
} else {
if (map[result[0]]!! >= value) {
// 堆顶 key 的出现次数大于当前 key, 不替换
return@forEach
}
var i = 0
var minIndex = 0
val maxIndex = result.lastIndex
// 用当前 key 替换堆顶 key, 并维护小顶堆性质
do {
result[i] = result[minIndex]
result[minIndex] = key
i = minIndex
var tmp = 2 * i + 1
if (tmp > maxIndex) {
break
}
if (map[result[tmp]]!! < map[result[minIndex]]!!) {
minIndex = tmp
}
tmp += 1
if (tmp <= maxIndex && map[result[tmp]]!! < map[result[minIndex]]!!) {
minIndex = tmp
}
} while (i != minIndex)
}
}
return result
}
}