使用大根堆结构。假设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个。
import heapq
def topk(arr1, arr2,k):# 放进堆,出来一个,把最靠近他的两个再放进堆
n1=len(arr1)
n2=len(arr2)
i=n1-1
j=n2-1
hp=[]
res=[]
heapq.heappush(hp,[-arr1[i]-arr2[j],i,j])
for _ in range(k):
t,i,j=heapq.heappop(hp)
res.append(-t)
a=[-arr1[i-1]-arr2[j],i-1,j]
b=[-arr1[i]-arr2[j-1],i,j-1]
heapq.heappush(hp,a)
heapq.heappush(hp,b)
return res
arr1=[1,3,4,5,7]
arr2=[1,2,5,8,9]
k=5
print(topk(arr1,arr2,k))
tql
真帅