unreasonably effective

you can be sloppy, as long as you are rigorous

LeetCode 373: Find K Pairs with Smallest Sums

link

For a pair (u, v), u \in \texttt{nums1} and v \in \texttt{nums2}.

Sweep line

We can group the pairs by distinct u‘s. Since a group has a fixed u and nums2 is sorted, we can consider each group as a sorted list of pairs (u, v). Keeping group heads in a min-heap, we can generate the pairs in sorted order of their sums with sweep line.

Say len(nums1) = n.

Time: \mathcal{O}(n + k \cdot \lg{n}), space: \mathcal{O}(n).

from heapq import heapify, heappush as push, heappop as pop


class Solution:
    def kSmallestPairs(
        self, nums1: List[int], nums2: List[int], k: int
    ) -> List[List[int]]:

        min_heap = [(first + nums2[0], first, nums2[0], 1) for first in nums1]
        heapify(min_heap)

        pairs = []
        while min_heap and len(pairs) < k:
            _, first, second, next_second_index = pop(min_heap)
            pairs.append([first, second])
            if next_second_index == len(nums2):
                continue

            second = nums2[next_second_index]
            push(min_heap, (first + second, first, second, next_second_index + 1))

        return pairs

tanvirdotzaman

,

Leave a comment