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

来源:互联网 发布:闲鱼 淘宝二手 编辑:程序博客网 时间:2024/06/03 17:45

【题目】

  给定两个有序数组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.5class Heap:    def __init__(self, row, col, value):        self.row = row        self.col = col        self.value = valuedef 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
原创粉丝点击