[力扣上岸] 堆算法 - 高频题Top k问题,你会做了吗?

301 阅读3分钟

前言

相信面试经验多的朋友一定见过top k问题。虽然是一道中等题,但是掌握了堆算法后就很简单了。

题目

复习一下【堆】

第一个类是自己实现的一个array,第二个类是最大堆。

class Array(object):
    """
    Achieve an Array by Python list
    """
    def __init__(self, size = 32):
        self._size = size
        self._items = [None] * size

    def __getitem__(self, index):
        """
        Get items
        :param index: get a value by index
        :return: value
        """
        return self._items[index]

    def __setitem__(self, index, value):
        """
        set item
        :param index: giving a index you want to teset
        :param value: the value you want to set
        :return:
        """
        self._items[index] = value

    def __len__(self):
        """
        :return: the length of array
        """
        return self._size

    def clear(self, value=None):
        """
        clear the Array
        :param value: set all value to None
        :return: None
        """
        for i in range(self._size):
            self._items[i] = value

    def __iter__(self):
        for item in self._items:
            yield item
            
class MaxHeap(object):
    def __init__(self, maxsize=None):
        self.maxsize = maxsize
        self._elements = Array(maxsize)
        self._count = 0

    def __len__(self):
        return self._count

    def add(self, value):
        if self._count >= self.maxsize:
            raise Exception('full')
        self._elements[self._count] = value
        self._count += 1
        self._siftup(self._count-1) 

    def _siftup(self, ndx):
        if ndx > 0:
            parent = int((ndx-1)/2)
            if self._elements[ndx] > self._elements[parent]: 
                self._elements[ndx], self._elements[parent] = self._elements[parent], self._elements[ndx]
                self._siftup(parent) 

    def extract(self):
        if self._count <= 0:
            raise Exception('empty')
        value = self._elements[0]
        self._count -= 1
        self._elements[0] = self._elements[self._count] 
        self._siftdown(0)
        return value

    def _siftdown(self, ndx):
        left = 2 * ndx + 1
        right = 2 * ndx + 2
        largest = ndx
        if (left < self._count and  
                self._elements[left] >= self._elements[largest] and
                self._elements[left] >= self._elements[right]): 
            largest = left
        elif right < self._count and self._elements[right] >= self._elements[largest]:
            largest = right
        if largest != ndx:
            self._elements[ndx], self._elements[largest] = self._elements[largest], self._elements[ndx]
            self._siftdown(largest)

梳理一下最大堆的四个重要函数:

add(value): 做了两件事:1. 把最新加入的value放在堆的最后面(也就是_count这个位置。_count统计的是数组最后一个位置)2.对此元素调用_siftup函数

siftup(ndx): 将第ndx元素的元素递归地与其父节点比大小,如果大于其父节点,则与其父节点交换位置,直到其小于其父节点或者没有父节点为止。

extract(): 两件事:1. 把数组第一个点(也就是最大堆的顶部)的值储存起来,最终要返回这个值。2. 把数组第一个点的值与最后一个值交换,并且让这个元素_siftdown。

siftdown(ndx): 将第ndx元素的元素递归地与其子节点比大小,如果小于其子节点,则与其子节点交换位置,直到其大于其子节点或者没有子节点为止。

重要性质:最大堆的父节点一定大于其所有子节点!再看一下刚才介绍的四个函数,其中,extract()可以把最大堆的顶部数值取出来,也就是当前堆的最大值。取k次就可以得到我们想要的top k数组。

堆排序(heapify)就是先构造出一个最大堆,然后不断extract(),把最大堆和最后一个元素交换位置,_count减一,直到堆变空,数组变成升序为止(因为每一次都把当前最大元素放到堆的最后面)。

题解

这里我使用了Python的heapq,就是Python已有的一个堆的库。

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = collections.defaultdict(int)
        for i in nums:
            count[i] += 1
        freqList = [(-freq, num) for (num, freq) in count.items()]
        heapq.heapify(freqList)
        res = [heapq.heappop(freqList)[1] for _ in range(k)]
        return res

先对数组中的所有元素统计一下出现次数,再把出现次数的相反数放进一个数组(因为heapq是最小堆,所有要用相反数来得到出现次数最多的元素)。用heapq对这个数组进行heapify排序,最终再取出排序完的前k个元素。

参考链接

[1]blog.csdn.net/qq_23869697…

[2]leetcode-cn.com/problems/to…

更多精彩请关注公众号:力扣上岸