Python heapq模块

这个模块(build-in)实现了一个堆的数据结构，完美的解决了Top-K问题，以后解决Top-K问题的时候，直接把这个模块拿来用就可以了

注意，默认的heap是一个小顶堆！

heapq模块提供了如下几个函数：

heapq.heappush(heap, item) 把item添加到heap中（heap是一个列表）

heapq.heappop(heap) 把堆顶元素弹出，返回的就是堆顶

heapq.heappushpop(heap, item) 先把item加入到堆中，然后再pop，比heappush()再heappop()要快得多

heapq.heapreplace(heap, item) 先pop，然后再把item加入到堆中，比heappop()再heappush()要快得多

heapq.heapify(x) 将列表x进行堆调整，默认的是小顶堆

heapq.merge(*iterables) 将多个列表合并，并进行堆调整，返回的是合并后的列表的迭代器

heapq.nlargest(n, iterable, key=None) 返回最大的n个元素（Top-K问题）

heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素（Top-K问题）

import heapq
import random
 
# Top-K
mylist = list(random.sample(range(100), 10))
k = 3
largest = heapq.nlargest(k, mylist)
smallest = heapq.nsmallest(k, mylist)
print('original list is', mylist)
print('largest-'+str(k), '  is ', largest)
print('smallest-'+str(k), ' is ', smallest)
 
# heapify
print('original list is', mylist)
heapq.heapify(mylist)
print('heapify  list is', mylist)
 
# heappush & heappop
heapq.heappush(mylist, 105)
print('pushed heap is', mylist)
heapq.heappop(mylist)
print('popped heap is', mylist)
 
# heappushpop & heapreplace
heapq.heappushpop(mylist, 130)    # heappush -> heappop
print('heappushpop', mylist)
heapq.heapreplace(mylist, 2)    # heappop -> heappush
print('heapreplace', mylist)

　输出结果

original list is [18, 6, 10, 24, 48, 2, 9, 25, 16, 89]
largest-3   is  [89, 48, 25]
smallest-3  is  [2, 6, 9]
original list is [18, 6, 10, 24, 48, 2, 9, 25, 16, 89]
heapify  list is [2, 6, 9, 16, 48, 10, 18, 25, 24, 89]
pushed heap is [2, 6, 9, 16, 48, 10, 18, 25, 24, 89, 105]
popped heap is [6, 16, 9, 24, 48, 10, 18, 25, 105, 89]
heappushpop [9, 16, 10, 24, 48, 130, 18, 25, 105, 89]
heapreplace [2, 16, 10, 24, 48, 130, 18, 25, 105, 89]

leetcode 703 python

设计一个找到数据流中第K大元素的类（class）。注意是排序后的第K大元素，不是第K个不同的元素。

你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。

示例:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

说明: 
你可以假设 nums 的长度≥ k-1 且k ≥ 1。

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.pool = nums
        self.size = len(self.pool)
        self.k = k
        heapq.heapify(self.pool)
        while self.size > k:
            heapq.heappop(self.pool)
            self.size -= 1

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if self.size < self.k:
            heapq.heappush(self.pool, val)
            self.size += 1
        elif val > self.pool[0]:
            heapq.heapreplace(self.pool, val)
        return self.pool[0]