【算法面试题python】合并k个有序列表

题目：合并k个有序list
输入：[[3,4,5], [1,6], [2,6,7,8], [5,6,9]]
输出：[1, 2, 3, 4, 5, 5, 6, 6, 6, 7, 8, 9]
思路：为每个list分配一个指针，每次循环比较k个指针指向的元素，取最小的放入新列表，并把对应的指针+1，退出循环的条件是所有列表的元素都被取完
代码：

def merge(a):
    new = []
    idx = [0]*len(a) # k个指针
    while True:
        curmin = 1e6  
        curidx = 0
        flag = 0
        for i in range(len(idx)):
            if idx[i]<len(a[i]):
                if a[i][idx[i]]< curmin:
                    curmin = a[i][idx[i]]
                    curidx = i # 最小值在第curidx个列表
            else:
                flag += 1
        if flag == len(idx):
            break
        new.append(curmin)
        idx[curidx] += 1 # 最小值所在的列表的指针要+1
    return new