[LeetCode] 347. Top K Frequent Elements（频率第 K 大的数）

• Difficulty: Medium

• Related Topics: Hash Table, Heap

• Link: https://leetcode.com/problems/top-k-frequent-elements/

Description

Given a non-empty array of integers, return the k most frequent elements.
给一个非空整数数组，返回频率第 k 大的数。

Examples

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.
  你可以假设 k 总是合理的（1 ≤ k ≤ 互异的元素个数）。

• Your algorithm's time complexity must be better than O(n log n), where n is the array's size.
  你的算法的时间复杂度必须优于 O(n log n), 其中 n 是数组的大小。

• It's guaranteed that the answer is unique, in other words the set of the top k freq elements is unique.
  输入保证答案只有一个，也就是说，频率第 k 大的数是唯一的。

• You can return the answer in any order.
  你可以以任意顺序返回答案。

Solution

一个朴素的想法是统计频率，然后将 (数, 频率) 对扔进堆里，再取出前 k 个，代码如下：

import java.util.*

class Solution {
    fun topKFrequent(nums: IntArray, k: Int): IntArray {
        val heap = PriorityQueue<Map.Entry<Int, Int>>(compareByDescending { it.value })
        heap.addAll(nums.countMap().entries)
        val result = arrayListOf<Int>()
        (1..k).forEach { result.add(heap.poll().key) }
        return result.toIntArray()
    }

    private fun IntArray.countMap(): Map<Int, Int> {
        val result = hashMapOf<Int, Int>()
        this.forEach { result[it] = result.getOrDefault(it, 0) + 1 }
        return result
    }
}

担心自己的时间复杂度不合要求，跑到 discussion 去看了一眼，发现了一个 O(N) 的解法，同样也是先统计频率，不过这次使用桶排来取代堆，就降低了时间复杂度。源码为 Java 实现，改写成 Kotlin 后代码如下：

class Solution {
    fun topKFrequent(nums: IntArray, k: Int): IntArray {
        val countMap = nums.countMap()

        val bucket = Array(nums.size + 1) { arrayListOf<Int>() }
        for ((num, freq) in countMap) {
            bucket[freq].add(num)
        }

        val result = arrayListOf<Int>()
        for (i in bucket.indices.reversed()) {
            if (result.size == k) {
                break
            }
            if (bucket[i].isNotEmpty()) {
                result.addAll(bucket[i])
            }
        }

        return result.toIntArray()
    }

    private fun IntArray.countMap(): Map<Int, Int> {
        val result = hashMapOf<Int, Int>()
        this.forEach { result[it] = result.getOrDefault(it, 0) + 1 }
        return result
    }
}

不过拿上去一跑，第一段代码的时间是 240ms，第二段代码的时间 268ms，果然 LeetCode 的计时器只能看看而已吗😂。