【题目】
给定两个有序数组arr1和arr2,再给定一个整数k,返回来自arr1和arr2的两个数相加和最大的前k个,两个数必须分别来自两个数组。要求时间复杂度O(klogk)。
【基本思路】
使用大根堆结构。假设arr1的长度是M,arr2的长度是N。因为是排序数组,arr1中最后一个数加上arr2中最后一个数一定就是最大的相加和。将这个数压入大根堆中。然后从大根堆中弹出一个堆顶,此时这个堆顶一定是(M-1, N-1)位置的和,表示获得一个最大相加和。然后,将两个相邻位置的和再放入堆中,即位置(M-1,N-2)和(M-2, N-1),因为除(M-1, N-1)位置的和外,最大的相加和一定在位置(M-1,N-2)和(m-2, N-1)中产生。重新调整大根堆,然后继续弹出,继续将弹出元素的两个相邻位置添加到堆中,直到弹出的元素达到K个。
【代码实现】
#python3.5
class Heap:
def __init__(self, row, col, value):
self.row = row
self.col = col
self.value = value
def getTopKSum(arr1, arr2, k):
def heapInsert(heap, row, col, data, i):
node = Heap(row, col, data)
heap[i] = node
parent = (i-1) // 2
while parent >= 0 and heap[parent].value < heap[i].value:
heap[parent], heap[i] = heap[i], heap[parent]
i = parent
parent = (i-1) // 2
def popHead(heap, heapSize):
res = heap[0]
heap[0], heap[heapSize-1] = heap[heapSize-1], heap[0]
heapify(heap, 0, heapSize-1)
return res
def heapify(heap, i, heapSize):
left = 2 * i + 1
right = 2 * i + 2
most = i
while left < heapSize:
if heap[left].value > heap[i].value:
most = left
if right < heapSize and heap[right].value > heap[most].value:
most = right
if most == i:
break
else:
heap[most], heap[i] = heap[i], heap[most]
i = most
left = 2 * i + 1
right = 2 * i + 2
def isContains(row, col, posSet):
return '_'.join([str(row),str(col)]) in posSet
def addPosToSet(row, col, posSet):
posSet.add('_'.join([str(row), str(col)]))
if arr1 == None or arr2 == None or k < 1 or k > len(arr1) * len(arr2):
return
heap = [0 for i in range(k+1)]
row = len(arr1) - 1
col = len(arr2) - 1
heapSize = 0
heapInsert(heap, row, col, arr1[row] + arr2[col], heapSize)
heapSize += 1
posSet = set()
count = 0
res = []
while count < k:
cur = popHead(heap, heapSize)
heapSize -= 1
res.append(cur.value)
r = cur.row
c = cur.col
if not isContains(r-1,c, posSet):
heapInsert(heap, r-1, c, arr1[r-1] + arr2[c], heapSize)
heapSize += 1
addPosToSet(r-1, c, posSet)
if not isContains(r, c-1, posSet):
heapInsert(heap, r, c-1, arr1[r] + arr2[c-1], heapSize)
heapSize += 1
addPosToSet(r, c-1, posSet)
count += 1
return res