其他题目---两个有序数组间相加和的TopK问题

【题目】

　　给定两个有序数组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