Problem statement

https://binarysearch.com/problems/K-Largest-Pairs

Solution

Similar to leetcode 373. Find K Pairs with Smallest Sums, one of the possible solutions is to use binary seach.

Complexity

Time complexity is O(n log n + m log m). Notice that this solution do not depend on k and will work upto k = m*n, so more difficult problem is solved in fact.

Code

from bisect import bisect, bisect_left

class Solution:
    def solve(self, A, B, k):
        if not A or not B: return 0

        A, B = sorted([-i for i in A]), sorted([-i for i in B])
        acc_B = [0] + list(accumulate(B))

        end, beg = A[-1] + B[-1], A[0] + B[0]
        while beg < end:
            mid = (beg + end) // 2
            j, cnt = len(B) - 1, 0
            for i in range(len(A)):
                while j >= 0 and A[i] + B[j] > mid:
                    j -= 1
                cnt += j + 1
            
            if cnt < k:
                beg = mid + 1
            else:
                end = mid

        ans, taken = 0, 0
        
        for x in A:
            idx = bisect(B, beg - x)
            taken += idx
            ans += acc_B[idx] + x*idx

        return -ans + (taken - k)*beg