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], index, i+1))
    return result

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

if __name__ == "__main__":
    test_heap_sort()