Skip to content

LeetCode 347.前 K 个高频元素

一、题目

[347] 前 K 个高频元素

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

示例 1:

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

示例 2:

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

说明:

你可以假设给定的 k 总是合理的,且 1 ≤ k ≤ 数组中不相同的元素的个数。 你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。

二、思路

可将这个问题分解为 2 个问题:

  1. 统计每个整数出现的次数;

  2. 获取前 k 个最大的值。

1. 解法一

第一步,确定每个整数出现的次数,可以简单地用一个字典对象维护每个值出现的次数。

遍历数组,每个数字出现时就给它的计数+1。

第二步,将每个整数出现的次数构建一个堆,依次取出堆的前 k 个元素。

python
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        自行实现堆结构
        """
        # 第一步:遍历数组,统计整数的出现次数
        # 时间复杂度:O(n)
        counts = {}
        for n in nums:
            counts[n] = counts.get(n, 0) + 1

        # 第二步:构建大根堆
        # 时间复杂度:O(nlogn)
        heap = BigRootHeap()
        for n, count in counts.items():
            node = Node(count, n)
            heap.push(node)

        # 第三步:依次取出前k个元素
        # 时间复杂度:O(logn)
        result = [heap.pop().value for _ in range(k)]
        return result

实现一个堆

堆是一个完全二叉树,堆的性质为:大根堆的每个节点都不小于其子节点,小根堆的每个节点都不大于其子节点。因此,大根堆的顶部必然最大的元素,而小根堆的顶部必然是最小的元素。

由于堆是完全的二叉树,可以用数组进行储存它的元素,只要知道了节点在列表中的下标就可以确定节点在堆中的位置。见下图:

img

img

如果节点的在数组中的下标值 rank,它的父节点下标为 (rank-1) // 2,左子节点(如果有)的下标为 rank * 2 + 1,右子节点的下标为 rank * 2 + 1。并且只有当 rank < n // 2 时(n 为数组元素个数),才有子节点。

建堆的过程,可以实现为向一个堆中逐渐添加元素(下述代码中的 push 方法)。

以大根堆为例:

  • 向一个已有的堆添加元素时,首先将元素添加到堆的末尾,再比较这个元素与它的父元素大小。
    1. 如果这个元素没有父元素(即这个元素是堆顶),或者这个元素不大于父元素,则所有元素都满足堆的性质。
    2. 如果这个元素大于父元素,则将这个元素与它的父元素交换,这样调整后这个局部就恢复了堆的性质。
    3. 接着再比较这个元素与它现在的父元素之间的大小,重复 1、2 步。最大的调整的次数与当前堆的高度有关,最多为 logn 次。

而从大根堆中取出最大元素的过程(下述代码中的 pop 方法),则可以实现为以下步骤:

  1. 取出堆顶元素,这个元素是所有元素中最大的。
  2. 将堆的最后一个元素放到堆顶,堆的元素总数减一。
  3. 从新的堆顶开始,判断元素与它的左右子节点的大小关系。
    • 如果这个元素没有子节点或不小于它的子节点,则整个堆都满足了堆的性质,不需要调整。
    • 否则,将这个元素与它最大的子节点交换,这样局部就恢复了堆的性质。接着再比较这个元素与它现在的子节点之间的大小关系,重复第 3 步。最大的调整的次数与当前堆的高度有关,最多为 logn 次。

代码实现如下:

python
class Node:
    """堆节点对象"""
    __slots__ = ('value', '_order')

    def __init__(self, order, value):
        self._order = order
        self.value = value

    def __gt__(self, other):
        """定义 > 操作符"""
        return self._order > other._order

    def __lt__(self, other):
        """定义 < 操作符"""
        return self._order < other._order


class BigRootHeap:
    """基于数组实现的大根堆"""
    def __init__(self):
        self._nodes = []

    def __len__(self):
        return len(self._nodes)

    @property
    def root(self):
        if self._nodes:
            return self._nodes[0]
        else:
            return None

    @staticmethod
    def get_parent_rank(rank: int):
        """获取父节点"""
        if rank > 0:
            return (rank - 1) // 2
        else:
            return None

    @staticmethod
    def get_left_child_rank(rank: int):
        """获取左子节点"""
        return rank * 2 + 1

    @staticmethod
    def get_right_child_rank(rank: int):
        """获取右子节点"""
        return rank * 2 + 2

    def push(self, node: Node):
        """向堆添加节点"""
        if not self._nodes:
            self._nodes = [node]
        else:
            # 添加到堆的末尾
            self._nodes.append(node)
            # 添加新元素后调整堆,以维持堆的性质
            self.heapify_up(len(self._nodes) - 1)

    def heapify_up(self, rank: int):
        """检查节点与其父节点是否满足大根堆的性质,如果不满足则调整"""
        if rank == 0:
            # 已经是根节点,不用调整
            return
        else:
            parent_rank = self.get_parent_rank(rank)
            if self._nodes[rank] > self._nodes[parent_rank]:
                # 将节点与父节点交换
                self._nodes[rank], self._nodes[parent_rank] = \
                    self._nodes[parent_rank], self._nodes[rank]
                # 递归对祖先节点进行调整,最大调整次数:logn
                self.heapify_up(parent_rank)
            else:
                # 当前节点与其父节点的关系不需要调整
                return

    def pop(self):
        """弹出当前堆顶元素"""
        if self._nodes:
            # 将末位元素交换到堆顶,取出最后一个元素
            self._nodes[0], self._nodes[-1] = self._nodes[-1], self._nodes[0]
            res = self._nodes.pop()
            # 从堆顶开始调整,维持堆性质
            self.heapify_down(0)
            return res

    def heapify_down(self, rank: int):
        """检查节点与其子节点是否满足大根堆的性质,如果不满足则调整"""
        if rank >= len(self._nodes) // 2:
            # 当前节点已经没有子节点了
            return
        else:
            lc_child_rank = self.get_left_child_rank(rank)
            # 如果当前节点比左子节点或右子节点小,则将其与最大的子节点交换,然后到相应的子节点位置继续调整。最大调整次数:logn
            biggest_rank = rank
            if self._nodes[biggest_rank] < self._nodes[lc_child_rank]:
                biggest_rank = lc_child_rank
            rc_child_rank = self.get_right_child_rank(rank)
            if rc_child_rank < len(self._nodes) and self._nodes[biggest_rank] < self._nodes[rc_child_rank]:
                biggest_rank = rc_child_rank

            if biggest_rank != rank:
                self._nodes[rank], self._nodes[biggest_rank] = \
                        self._nodes[biggest_rank], self._nodes[rank]
                self.heapify_down(biggest_rank)
            else:
                return

2. 解法二(优化堆空间)

由于题目只需要出现频率最高的前 k 个元素,用大根堆,需要把所有元素都入堆,而如果用小根堆,则只需要将最大的 k 个元素入堆,其余元素可以抛弃。

python
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        改进:使用小根堆,仅将最小的k个元素入堆,减小空间
        """
        # 第一步:遍历数组,统计整数的出现次数
        # 时间复杂度:O(n)
        counts = {}
        for n in nums:
            counts[n] = counts.get(n, 0) + 1

        # 第二步:构建小根堆
        # 时间复杂度:O(nlogn)
        # 堆的空间复杂度为常数级
        heap = SmallRootHeap()
        for n, count in counts.items():
            node = Node(count, n)
            if len(heap) >= k:
                if heap.root > node:
                    # 堆已经有k个元素,且新元素比堆中的最小元素小,则不需要入堆
                    continue
                else:
                    # 堆中最小元素出堆
                    heap.pop()
            heap.push(node)

        # 第三步:依次取出前k个元素,然后逆序
        # 时间复杂度:O(n)
        result = [heap.pop().value for _ in range(k)]
        result.reverse()
        return result

3. 利用 python 自带的堆结构

python 的标准库中已经包含了堆结构的实现,导入 heapq 这个包就可以使用。实际工作中有需要使用堆结构(如,优先级队列)时可以利用这个库。

python
import heapq

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        用python自带的堆结构
        """
        # 第一步:遍历数组,统计整数的出现次数
        counts = {}
        for n in nums:
            counts[n] = counts.get(n, 0) + 1
        nodes = (Node(count, num) for num, count in counts.items())

        # 第二步,用python自带的堆结构完成建堆和取前k个值
        n_largest = heapq.nlargest(k, nodes)
        return [x.value for x in n_largest]

4. 只要一行代码的终极偷懒法

python 标准库的 collection 包有一个 Counter 类,可以接受任意的可哈希对象序列作为输入,调用它的 most_common()方法即可获得结果。

当实际遇到需要对数据进行计数的需求时,用这个方法的开发效率和运行效率都是最高的。

python
from collections import Counter

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        用python自带的Counter类
        """
        return [x[0] for x in Counter(nums).most_common(k)]

5. 空间换时间:桶排序

传统的排序算法都无法超越 O(nlogn)的时间复杂度,但在特殊情况下,可以做到比 O(nlogn)更快。

这道题的我们已经确定了有 n 个整数,可以确定每个整数的出现次数的所有可能取值的范围区间是 [0, n]。因此,可以分配一个大小为 n+1 的数组(这个数组即所谓的“桶”),数组下标为 i 的值为出现次数为 i 次的所有整数。

我们先遍历一遍整数,统计每个值的出现次数。然后再遍历一次,将每个元素放到“桶”的相应位置。最后逆序取出 k 个元素即可。

这个算法的整体的时间复杂度为 O(n),空间复杂度也为 O(n)。

python
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """桶排序方法"""
        # 第一步:遍历数组,统计整数的出现次数
        # 时间复杂度:O(n)
        counts = {}
        for n in nums:
            counts[n] = counts.get(n, 0) + 1

        # 第二步:桶排序
        # 时间复杂度:O(n)
        bucket = [[]]*len(nums)
        for num, count in counts.items():
            if bucket[count]:
                bucket[count].append(num)
            else:
                bucket[count] = [num]

        # 第三步:取出桶的末尾k个元素
        res = []
        for x in reversed(bucket):
            if k <= 0:
                break
            for num in x:
                res.append(num)
                k -= 1
                if k <= 0:
                    break
        return res

Released under the MIT License.