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.

思路

分成两步来解决:

1. 先用 Map 统计下每个数出现的次数, 剩下的就是一个 TopN 问题;
2. 遍历 Map 中的键值对, 遍历过程中维护一个容量为K的小顶堆.

小顶堆维护过程如下:

1. 如果小顶堆中元素不足K个, 将当前元素添加到堆中;
2. 如果小顶堆中有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
    }
}