373 find k pairs with smallest sums

step1-1.py

@@ -0,0 +1,23 @@
+from typing import List
+import heapq
+
+class Solution:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        already_in_heap = set()
+        candidates_heap = []
+        result = []
+        num_taken = 0
+        if len(nums1) * len(nums2) < k:
+            raise ValueError("Given arrays wouldn't yield enough pairs.")
+        heapq.heappush(candidates_heap, (nums1[0] + nums2[0], 0, 0))
+        while num_taken < k:
+            _, smallest_idx1, smallest_idx2 = heapq.heappop(candidates_heap)
+            result.append((nums1[smallest_idx1], nums2[smallest_idx2]))
+            num_taken += 1
+            if smallest_idx1 < len(nums1) - 1 and (smallest_idx1 + 1, smallest_idx2) not in already_in_heap:
+                heapq.heappush(candidates_heap, (nums1[smallest_idx1 + 1] + nums2[smallest_idx2], smallest_idx1 + 1, smallest_idx2))
+                already_in_heap.add((smallest_idx1 + 1, smallest_idx2))
+            if smallest_idx2 < len(nums2) - 1 and (smallest_idx1, smallest_idx2 + 1) not in already_in_heap:
+                heapq.heappush(candidates_heap, (nums1[smallest_idx1] + nums2[smallest_idx2 + 1], smallest_idx1, smallest_idx2 + 1))
+                already_in_heap.add((smallest_idx1, smallest_idx2 + 1))
+        return result

step2-1.py

@@ -0,0 +1,25 @@
+from typing import List
+import heapq
+
+class Solution:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        seen_idx1_idx2_pairs = set()
+        sum_idx1_idx2_tuple_heap = []
+        result = []
+        num_taken = 0
+        if len(nums1) * len(nums2) < k:
+            raise ValueError("Given arrays wouldn't yield enough pairs.")
+        heapq.heappush(sum_idx1_idx2_tuple_heap, (nums1[0] + nums2[0], 0, 0))
+        while num_taken < k:
+            _, cur_smallest_idx1, cur_smallest_idx2 = heapq.heappop(sum_idx1_idx2_tuple_heap)
+            result.append((nums1[cur_smallest_idx1], nums2[cur_smallest_idx2]))
+            num_taken += 1
+            next_idx_of_idx1 = cur_smallest_idx1 + 1
+            next_idx_of_idx2 = cur_smallest_idx2 + 1
+            if next_idx_of_idx1 < len(nums1) and (next_idx_of_idx1, cur_smallest_idx2) not in seen_idx1_idx2_pairs:
+                heapq.heappush(sum_idx1_idx2_tuple_heap, (nums1[next_idx_of_idx1] + nums2[cur_smallest_idx2], next_idx_of_idx1, cur_smallest_idx2))
+                seen_idx1_idx2_pairs.add((next_idx_of_idx1, cur_smallest_idx2))
+            if next_idx_of_idx2 < len(nums2) and (cur_smallest_idx1, next_idx_of_idx2) not in seen_idx1_idx2_pairs:
+                heapq.heappush(sum_idx1_idx2_tuple_heap, (nums1[cur_smallest_idx1] + nums2[next_idx_of_idx2], cur_smallest_idx1, next_idx_of_idx2))
+                seen_idx1_idx2_pairs.add((cur_smallest_idx1, next_idx_of_idx2))
+        return result

step2-2.py

@@ -0,0 +1,53 @@
+from typing import List
+import heapq
+
+
+class Solution:
+    def _append_sum_idx_tuple_if_possible(self, idx1, idx2, candidate_heap, taken_count_from_nums1_list, taken_count_from_nums2_list, nums1, nums2):
+        if idx1 < 0 or idx1 >= len(nums1):
+            return
+        if idx2 < 0 or idx2 >= len(nums2):
+            return
+        is_idx1_takable = taken_count_from_nums1_list[idx1] == idx2
+        is_idx2_takable = taken_count_from_nums2_list[idx2] == idx1
+        if is_idx1_takable and is_idx2_takable:
+            heapq.heappush(candidate_heap, (nums1[idx1] + nums2[idx2], idx1, idx2))
+            return
+        return
+
+
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        nums1_len = len(nums1)
+        nums2_len = len(nums2)
+        taken_count_from_nums1_list = [0] * nums1_len
+        taken_count_from_nums2_list = [0] * nums2_len
+        candidate_sum_idx1_idx2_tuple = []
+        heapq.heappush(candidate_sum_idx1_idx2_tuple, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if len(candidate_sum_idx1_idx2_tuple) == 0:
+                print("Not enough elements in nums1 and nums2. Try after changing k.")
+                return result
+            _, idx1, idx2 = heapq.heappop(candidate_sum_idx1_idx2_tuple)
+            result.append((nums1[idx1], nums2[idx2]))
+            taken_count_from_nums1_list[idx1] += 1
+            taken_count_from_nums2_list[idx2] += 1
+            self._append_sum_idx_tuple_if_possible(
+                idx1 + 1,
+                idx2,
+                candidate_sum_idx1_idx2_tuple,
+                taken_count_from_nums1_list,
+                taken_count_from_nums2_list,
+                nums1,
+                nums2
+            )
+            self._append_sum_idx_tuple_if_possible(
+                idx1,
+                idx2 + 1,
+                candidate_sum_idx1_idx2_tuple,
+                taken_count_from_nums1_list,
+                taken_count_from_nums2_list,
+                nums1,
+                nums2
+            )
+        return result

step2-3.py

@@ -0,0 +1,32 @@
+from typing import List
+import heapq
+
+class Solution:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        taken_count1 = [0] * len(nums1)
+        taken_count2 = [0] * len(nums2)
+        sum_idx1_idx2_heap = []
+        def append_if_possible(idx1, idx2):
+            if idx1 < 0 or idx1 >= len(nums1):
+                return
+            if idx2 < 0 or idx2 >= len(nums2):
+                return
+            is_idx1_takable = taken_count1[idx1] == idx2
+            is_idx2_takable = taken_count2[idx2] == idx1
+            if is_idx1_takable and is_idx2_takable:
+                heapq.heappush(sum_idx1_idx2_heap, (nums1[idx1] + nums2[idx2], idx1, idx2))
+            return
+
+        heapq.heappush(sum_idx1_idx2_heap, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if len(sum_idx1_idx2_heap) == 0:
+                print("Not enough elements in nums1 and nums2. Try after changing k.")
+                return
+            _, idx1, idx2 = heapq.heappop(sum_idx1_idx2_heap)
+            taken_count1[idx1] += 1
+            taken_count2[idx2] += 1
+            result.append((nums1[idx1], nums2[idx2]))
+            append_if_possible(idx1 + 1, idx2)
+            append_if_possible(idx1, idx2 + 1)
+        return result

step3-1.py

@@ -0,0 +1,21 @@
+from typing import List
+import heapq
+
+class Solution:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        already_in_heap_or_taken = set()
+        candidate = []
+        heapq.heappush(candidate, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if not candidate:
+                raise ValueError("k is too large")
+            _, idx1, idx2 = heapq.heappop(candidate)
+            result.append((nums1[idx1], nums2[idx2]))
+            if idx1 + 1 < len(nums1) and (idx1 + 1, idx2) not in already_in_heap_or_taken:
+                already_in_heap_or_taken.add((idx1 + 1, idx2))
+                heapq.heappush(candidate, (nums1[idx1 + 1] + nums2[idx2], idx1 + 1, idx2))
+            if idx2 + 1 < len(nums2) and (idx1, idx2 + 1) not in already_in_heap_or_taken:
+                already_in_heap_or_taken.add((idx1, idx2 + 1))
+                heapq.heappush(candidate, (nums1[idx1] + nums2[idx2 + 1], idx1, idx2 + 1))
+        return result

step3-3.py

@@ -0,0 +1,31 @@
+from typing import List
+import heapq
+
+class Solution:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        num_taken1 = [0] * len(nums1)
+        num_taken2 = [0] * len(nums2)
+        candidate = []
+        def append_if_possible(idx1, idx2):
+            if idx1 < 0 or idx1 >= len(nums1):
+                return
+            if idx2 < 0 or idx2 >= len(nums2):
+                return
+            is_idx1_takable = num_taken1[idx1] == idx2
+            is_idx2_takable = num_taken2[idx2] == idx1
+            if is_idx1_takable and is_idx2_takable:
+                heapq.heappush(candidate, (nums1[idx1] + nums2[idx2], idx1, idx2))
+            return
+
+        heapq.heappush(candidate, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if not candidate:
+                raise ValueError("k is too large")
+            _, idx1, idx2 = heapq.heappop(candidate)
+            result.append((nums1[idx1], nums2[idx2]))
+            num_taken1[idx1] += 1
+            num_taken2[idx2] += 1
+            append_if_possible(idx1 + 1, idx2)
+            append_if_possible(idx1, idx2 + 1)
+        return result

time_comparison.py

@@ -0,0 +1,121 @@
+import timeit
+import random
+from typing import List
+import heapq
+
+
+class SolutionClassMethod:
+    def _append_sum_idx_tuple_if_possible(self, idx1, idx2, candidate_heap, taken_count_from_nums1_list, taken_count_from_nums2_list, nums1, nums2):
+        if idx1 < 0 or idx1 >= len(nums1):
+            return
+        if idx2 < 0 or idx2 >= len(nums2):
+            return
+        is_idx1_takable = taken_count_from_nums1_list[idx1] == idx2
+        is_idx2_takable = taken_count_from_nums2_list[idx2] == idx1
+        if is_idx1_takable and is_idx2_takable:
+            heapq.heappush(candidate_heap, (nums1[idx1] + nums2[idx2], idx1, idx2))
+            return
+        return
+
+
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        nums1_len = len(nums1)
+        nums2_len = len(nums2)
+        taken_count_from_nums1_list = [0] * nums1_len
+        taken_count_from_nums2_list = [0] * nums2_len
+        candidate_sum_idx1_idx2_tuple = []
+        heapq.heappush(candidate_sum_idx1_idx2_tuple, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if len(candidate_sum_idx1_idx2_tuple) == 0:
+                print("Not enough elements in nums1 and nums2. Try after changing k.")
+                return result
+            _, idx1, idx2 = heapq.heappop(candidate_sum_idx1_idx2_tuple)
+            result.append((nums1[idx1], nums2[idx2]))
+            taken_count_from_nums1_list[idx1] += 1
+            taken_count_from_nums2_list[idx2] += 1
+            self._append_sum_idx_tuple_if_possible(
+                idx1 + 1,
+                idx2,
+                candidate_sum_idx1_idx2_tuple,
+                taken_count_from_nums1_list,
+                taken_count_from_nums2_list,
+                nums1,
+                nums2
+            )
+            self._append_sum_idx_tuple_if_possible(
+                idx1,
+                idx2 + 1,
+                candidate_sum_idx1_idx2_tuple,
+                taken_count_from_nums1_list,
+                taken_count_from_nums2_list,
+                nums1,
+                nums2
+            )
+        return result
+
+
+class SolutionInnerMethod:
+    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
+        taken_count1 = [0] * len(nums1)
+        taken_count2 = [0] * len(nums2)
+        sum_idx1_idx2_heap = []
+        def append_if_possible(idx1, idx2):
+            if idx1 < 0 or idx1 >= len(nums1):
+                return
+            if idx2 < 0 or idx2 >= len(nums2):
+                return
+            is_idx1_takable = taken_count1[idx1] == idx2
+            is_idx2_takable = taken_count2[idx2] == idx1
+            if is_idx1_takable and is_idx2_takable:
+                heapq.heappush(sum_idx1_idx2_heap, (nums1[idx1] + nums2[idx2], idx1, idx2))
+            return
+
+        heapq.heappush(sum_idx1_idx2_heap, (nums1[0] + nums2[0], 0, 0))
+        result = []
+        while len(result) < k:
+            if len(sum_idx1_idx2_heap) == 0:
+                print("Not enough elements in nums1 and nums2. Try after changing k.")
+                return
+            _, idx1, idx2 = heapq.heappop(sum_idx1_idx2_heap)
+            taken_count1[idx1] += 1
+            taken_count2[idx2] += 1
+            result.append((nums1[idx1], nums2[idx2]))
+            append_if_possible(idx1 + 1, idx2)
+            append_if_possible(idx1, idx2 + 1)
+        return result
+
+def benchmark():
+    nums1 = list(range(1000000))
+    nums2 = list(range(1000000))
+    k = 10000
+
+    sol_class_method = SolutionClassMethod()
+    sol_inner_method = SolutionInnerMethod()
+
+    t_method = timeit.timeit(
+        stmt="sol_class_method.kSmallestPairs(nums1, nums2, k)",
+        globals=locals(),
+        number=50
+    )
+
+    t_inner = timeit.timeit(
+        stmt="sol_inner_method.kSmallestPairs(nums1, nums2, k)",
+        globals=locals(),
+        number=50
+    )
+    return t_method, t_inner
+
+
+if __name__ == "__main__":
+    method_times = []
+    inner_times = []
+    for _ in range(100):
+        t_method, t_inner = benchmark()
+        method_times.append(t_method)
+        inner_times.append(t_inner)
+    method_time = sum(method_times) / len(method_times)
+    inner_time = sum(inner_times) / len(inner_times)
+    print(f"Class Method: {method_time}")
+    print(f"Inner Method: {inner_time}")
+