heapq 对有序的数组列表进行整体排序

"""
功能：实现对有序的多个数组整体排序，获取top k个最小元素
"""
from heapq import *


def heap_sort(arr,top_k):
    q = []
    for i in range(len(arr)):
        heappush(q,(arr[i][0],i,0))
    result = []
    for k in range(top_k):
        if q:
            data,index,i = heappop(q)
            result.append(data)
            if i+1 < len(arr[index]):
                heappush(q,(arr[index][i+1],i+1))
    return result

def test_heap_sort():
    print(heap_sort(arr=[[1,3,4,9,80,799],[0,6,8,12,999],[123,345,987]],top_k=10))
    print(heap_sort(arr=[[1,top_k=3))
    print(heap_sort(arr=[[1,799]],top_k=0))


if __name__ == "__main__":
    test_heap_sort()