合并k个有序数组

Posted 2021-01-15 tanshoudong

给定K个有序数组，每个数组有n个元素，想把这些数组合并成一个有序数组

可以利用最小堆完成，时间复杂度是O(nklogk)，具体过程如下：

创建一个大小为n*k的数组保存最后的结果
创建一个大小为k的最小堆，堆中元素为k个数组中的每个数组的第一个元素
重复下列步骤n*k次：
每次从堆中取出最小元素（堆顶元素），并将其存入输出数组中
用堆顶元素所在数组的下一元素将堆顶元素替换掉，
如果数组中元素被取光了，将堆顶元素替换为无穷大。每次替换堆顶元素后，重新调整堆
初始化最小堆的时间复杂度O（k），总共有kn次循环，每次循环调整最小堆的时间复杂度是O（logk），所以总的时间复杂度是O（knlogk）

代码如下，已测试通过：

import sys

class HeapNode:
    def __init__(self,x,y=0,z=0):
        self.value=x
        self.i=y
        self.j=z


def Min_Heap(heap):#构造一个堆，将堆中所有数据重新排序
    HeapSize = len(heap)#将堆的长度单独拿出来方便
    for i in range((HeapSize -2)//2,-1,-1):#从后往前出数
        Min_Heapify(heap,i)


def Min_Heapify(heap,root):
    heapsize=len(heap)
    MIN=root
    left=2*root+1
    right=left+1
    if left<heapsize and heap[MIN].value>heap[left].value:
        MIN=left
    if right <heapsize and heap[MIN].value>heap[right].value:
        MIN=right
    if MIN!=root:
        heap[MIN], heap[root] = heap[root], heap[MIN]
        Min_Heapify(heap, MIN)

def MergeKArray(nums,n):
    # 合并k个有序数组，每个数组长度都为k
    knums=[]
    output=[]
    for i in range(len(nums)):
        subnums=nums[i]
        knums.append(HeapNode(subnums[0],i,1))
        Min_Heap(knums)#k个元素初始化最小堆

    for i in range(len(nums)*n):
        # 取堆顶，存结果
        root=knums[0]
        output.append(root.value)
        #替换堆顶
        if root.j<n:
            root.value=nums[root.i][root.j]
            root.j+=1
        else:
            root.value=sys.maxsize
        knums[0]=root
        Min_Heapify(knums,0)
    return output


knums=[[1,2,3],[1,3,6],[4,5,8]]
print(MergeKArray(knums,3))