Find K Pairs with Smallest Sums

Data Structures and Algorithms

Statement

Given two lists and an integer k, find k pairs of numbers with the smallest sum. Each list should contribute a number to each pair.

Constraints

• 1 ≤ list1.length(), list2.length() ≤ 500
• -10⁴ ≤ list1[i], list2[i] ≤ 10⁴
• 1 ≤ k ≤ 10³
• Input lists are sorted in ascending order
• If k exceeds the total number of valid pairs, then return all pairs

Solution

"""
production algorithm
"""

from data_structures.heap.heap import Heap


class Solution:
    def k_smallest_pairs(self, list1, list2, k):
        """
        time complexity O(klogn)
        space complexity O(n)
        """
        # add to heap sum of each value of first list
        # and first value of second list
        smallest = Heap()
        [smallest.push(MinHeapData(list1[i] + list2[0], i, 0))
         for i in range(len(list1))]
        result = list()

        for _ in range(k):
            if smallest.empty():
                break

            # add values of first and second list
            # to result seen so far
            _, index1, index2 = smallest.pop()
            result.append((list1[index1], list2[index2]))

            # add to heap sum of value of first list
            # and next value of second list
            if index2 + 1 < len(list2):
                smallest.push(MinHeapData(
                    list1[index1] + list2[index2 + 1], index1, index2 + 1))

        return result


class MinHeapData:
    def __init__(self, sum, index1, index2):
        self.sum = sum
        self.index1 = index1
        self.index2 = index2

    def __lt__(self, other):
        return self.sum < other.sum

    def __iter__(self):
        yield self.sum
        yield self.index1
        yield self.index2

    def __len__(self):
        return 3

"""
unit tests
"""

from unittest import TestCase
from algorithms.k_way_merge.find_k_pairs_with_smallest_sums.solution import Solution


class SolutionTestCase(TestCase):
    def test_average_case(self):
        # Given
        list1 = [0, 4, 4, 5, 16, 17, 18]
        list2 = [3, 7, 8, 9, 16, 17, 19]
        k = 5
        solution = Solution()

        # When
        actual = solution.k_smallest_pairs(list1, list2, k)

        # Then
        expected = [(0, 3), (4, 3), (0, 7), (4, 3), (0, 8)]
        self.assertEqual(actual, expected)

    def test_index_out_of_range_prevented(self):
        # Given
        list1 = [0, 104, 104, 105, 116, 117, 118]
        list2 = [3, 7, 8, 9, 16, 17, 19]
        k = 8
        solution = Solution()

        # When
        actual = solution.k_smallest_pairs(list1, list2, k)

        # Then
        expected = [(0, 3), (0, 7), (0, 8), (0, 9),
                    (0, 16), (0, 17), (0, 19), (104, 3)]
        self.assertEqual(actual, expected)

    def test_number_of_pairs_less_than_k(self):
        # Given
        list1 = [0, 4, 4]
        list2 = [3, 7, 8]
        k = 100
        solution = Solution()

        # When
        actual = solution.k_smallest_pairs(list1, list2, k)

        # Then
        expected = [(0, 3), (4, 3), (4, 3), (0, 7), (0, 8),
                    (4, 7), (4, 7), (4, 8), (4, 8)]
        self.assertEqual(actual, expected)

Strategy

K-way Merge.

Explanation

First, iterate list1. Add the sum of each number of list1 and the first number of list2 to a min heap. The size of the min heap will be the size of list1.

Next, iterate at most k times. For each iteration, pop the smallest sum from the min heap and add its pair of numbers to the result seen so far. Then, add then sum of the number of the first list and the next number of the second list to the min heap.

Finally, the result is the k pairs of numbers with the smallest sum.

Time Complexity

The time complexity of the algorithm depends on the size of the min heap and the number of operations to the min heap. The size of the min heap is the length of list1, and the number of operations to the min heap is k. Therefore, the time complexity of the algorithm is O(klogn).

Space Complexity

The auxiliary space complexity depends on the size of the min heap. Therefore, the auxiliary space complexity of the algorithm is O(n).

Links

data_structures_algorithms/algorithms/k_way_merge/find_k_pairs_with_smallest_sums at main ·…