373.Find K Pairs with Smallest Sums

Tags: [heap]

Link: https://leetcode.com/problems/find-k-pairs-with-smallest-sums/#/description

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair(u,v)which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1), (u2,v2) ... (uk,vk) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

Solution: Heap

class Solution(object):

    class HeapNode(object):
        def __init__(self, elem1, elem2, elem2_index):
            self.elem1 = elem1
            self.elem2 = elem2
            self.elem2_index = elem2_index

        def __cmp__(self, other):
            return (self.elem1 + self.elem2) - (other.elem1 + other.elem2)

    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        if not nums1 or not nums2 or not k:
            return []

        min_heap = []
        for elem1 in nums1:
            min_heap.append(self.HeapNode(elem1, nums2[0], 0))

        nums1_len = len(nums1)
        nums2_len = len(nums2)
        result = []
        for _ in xrange(min(k, nums1_len * nums2_len)):
            heap_root = heapq.heappop(min_heap)
            result.append([heap_root.elem1, heap_root.elem2])

            if heap_root.elem2_index + 1 < nums2_len:
                heapq.heappush(min_heap, self.HeapNode(heap_root.elem1, 
                                                       nums2[heap_root.elem2_index + 1], 
                                                       heap_root.elem2_index + 1))
        return result

Revelation:

• Next time, when we meet the problems which is asking for the kth smallest / biggest pairs, and there are multiple dimensional inputs, and each dimensional input is sorted, we can think about using the heap in this way to solve these problems.
• We use out of all elements from one dimension, and each time heap pop one node, then heap push another dimension element into heap.

Note:

• Time complexity = O(nums1_len + min(k, nums1_len * nums2_len * lg (nums1_len * nums2_len))), nums1_len is the number of elements in nums1, and nums2_len is the number of elements in nums2.