面试百度的时候遇到的一道题,输入为多个有序数组,输出为合并的一个有序数组。
思路:与“两个有序数组合并为一个有序数组”这道题相比,因为输入为多个有序数组(这里设有K个数组),因此每次都需要比较K个数的大小,取出最小的数存入返回的数组中,而对多个数取最小的操作,使用最小堆效率最高,因此这里就是 主要考察的数组结构就是最小堆。
因为python的堆结构需要自行设计,网上也没有看到比较好的回答,故自己整理如下。
输入:arr, n, m
其中arr的大小为[n,m], n为数组的个数,m为每个数组的长度
class node(object):
def __init__(self, value, k, j):
self.value = value
self.out_index = k
self.in_index = j
class solution():
def MergeArr(self, arr, n, m):
nums = []
for i in range(n):
nums.append(node(arr[i][0], i, 0))
for i in range(n):
self.heapInsert(nums, i)
res = []
for i in range(n*m):
value = nums[0].value
res.append(value)
k = nums[0].out_index
j = nums[0].in_index
if j+1<m:
nums[0] = node(arr[k][j+1], k, j+1)
else:
nums[0] = node(float('INF'), k, j+1)
self.heapify(nums)
return res
def heapInsert(self, nums, i):
'''
构建堆,自下而上
'''
while i>0 and nums[i].value<nums[(i-1)//2].value:
nums[i],nums[(i-1)//2] = nums[(i-1)//2], nums[i]
i = (i-1)//2
def heapify(self, nums):
'''
重建堆,自上而下
'''
n = len(nums)
index = 0
left = 2*index+1
while left<n:
if left+1<n and nums[left+1].value<nums[left].value:
smallest = left+1
else:
smallest = left
if nums[smallest].value<nums[index].value:
nums[smallest], nums[index] = nums[index], nums[smallest]
index = smallest
left = 2*index+1
else:
break
可以用一个测试用例测试:
arr = [[1,4,5,17],[0,2,14,28],[1,3,6,19],[6,15,23,24]]
s = solution()
print(s.MergeArr(arr, 4, 4))