heapq实现小顶堆(TopK大)、大顶堆(BtmK小)

Apr 25, 2024
2 views
Algorithm
class TopkHeap(object):
    def __init__(self, k):
        self.k = k
        self.data = []

    def Push(self, elem):
        if len(self.data) < self.k:
            heapq.heappush(self.data, elem)
        else:
            topk_small = self.data[0]
            if elem > topk_small:
                heapq.heapreplace(self.data, elem)

    def TopK(self):
        #return [x for x in reversed([heapq.heappop(self.data) for _ in range(len(self.data))])]
                return reversed(sorted(self.data))

自己实现小顶堆

class MinHeap:
    def __init__(self, size):
        self.size = size
        self.data = []

    def pushpop(self, v):
        if len(self.data) < self.size:
            self.data.append(v)
            self.swim(len(self.data) - 1)
        elif v > self.data[0]:
            self.data[0] = v
            self.sink(0)

    def large(self, p, q):
        return self.data[p] > self.data[q]

    def exch(self, p, q):
        self.data[p], self.data[q] = self.data[q], self.data[p]

    def swim(self, node):
        while (node - 1) // 2 >= 0 and self.large((node - 1) // 2, node):
            self.exch((node - 1) // 2, node)
            node = (node - 1) // 2

    def sink(self, node):
        while node * 2 < self.size - 1:
            left = node * 2 + 1
            right = node * 2 + 2
            child_min = left
            if right < self.size and self.large(left, right):
                child_min = right
            if not self.large(node, child_min):
                break
            self.exch(node, child_min)
            node = child_min


class Solution:
    def getLargestNumbers(self, arr: List[int], k: int) -> List[int]:

        if not k:
            return []

        mq = MinHeap(k)
        for num in arr:
            mq.pushpop(num)
        return sorted(mq.data, reverse=True)

变态的需求来了:给出N长的序列,求出BtmK小的元素,即使用大顶堆。

概括一种最简单的:

push(e)改为push(-e)pop(e)改为-pop(e)

也就是说,在存入堆、从堆中取出的时候,都用相反数,而其他逻辑与TopK完全相同,看代码:

class BtmkHeap(object):
    def __init__(self, k):
        self.k = k
        self.data = []

    def Push(self, elem):
        # Reverse elem to convert to max-heap
        elem = -elem
        # Using heap algorighem
        if len(self.data) < self.k:
            heapq.heappush(self.data, elem)
        else:
            topk_small = self.data[0]
            if elem > topk_small:
                heapq.heapreplace(self.data, elem)

    def BtmK(self):
        return sorted([-x for x in self.data])

自己实现大顶堆

class MaxHeap:
    def __init__(self, size):
        self.size = size
        self.data = []

    def pushpop(self, v):
        if len(self.data) < self.size:
            self.data.append(v)
            self.swim(len(self.data) - 1)
        elif v < self.data[0]:
            self.data[0] = v
            self.sink(0)

    def less(self, p, q):
        return self.data[p] <= self.data[q]

    def exch(self, p, q):
        self.data[p], self.data[q] = self.data[q], self.data[p]

    def swim(self, node):
        while (node - 1) // 2 >= 0 and self.less((node - 1) // 2, node):
            self.exch((node - 1) // 2, node)
            node = (node - 1) // 2

    def sink(self, node):
        while node * 2 < self.size - 1:
            left = node * 2 + 1
            right = node * 2 + 2
            child_max = left
            if right < self.size and self.less(left, right):
                child_max = right
            if not self.less(node, child_max):
                break
            self.exch(node, child_max)
            node = child_max


class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:

        if not k:
            return []

        mq = MaxHeap(k)
        for num in arr:
            mq.pushpop(num)
        return sorted(mq.data)