K Max Sum Combinations

'''
After sorting both arrays, the largest possible sum is always formed by the last elements, so we start from that pair. From any chosen pair (i, j), the only next possible better candidates are (i-1, j) and (i, j-1), because decreasing either index gives the next largest sums. We use a max heap to always pick the current largest sum and a visited set to avoid repeating the same index pairs until we collect the top k sums.
'''

from heapq import heappush,heappop
class Solution:
    def topKSumPairs(self, a, b, k):
        a.sort()
        b.sort()

        i,j=len(a),len(b)
        visited=set()
        heap=[]

        visited.add((i-1,j-1))
        last_sum=a[i-1]+b[j-1]
        heappush(heap,(-last_sum,(i-1,j-1)))
        ans=[]

        while len(heap)>0 and len(ans)<k:
            temp=heappop(heap)
            val=-temp[0]
            ans.append(val)
            i,j=temp[1][0],temp[1][1]

            if (i-1,j) not in visited:
                cur_sum=a[i-1]+b[j]
                heappush(heap,(-cur_sum,(i-1,j)))
                visited.add((i-1,j))
            if (i,j-1) not in visited:
                cur_sum=a[i]+b[j-1]
                heappush(heap,(-cur_sum,(i,j-1)))
                visited.add((i,j-1))
        return ans