LeetCode: Find K Pairs with Smallest Sums

Posted on August 18, 2019July 26, 2020 by braindenny

Find K Pairs with Smallest Sums

Similar Problems:

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:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: 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:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: 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:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Github: code.dennyzhang.com

Credits To: leetcode.com

Leave me comments, if you have better ways to solve.

Solution:

## https://code.dennyzhang.com/find-k-pairs-with-smallest-sums
## Basic Ideas: heap
##
##  Only maintain k elements in the heap
##
## Complexity: Time O(k*log(k)), Space O(k)
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        nums = []
        for v1 in nums1:
            for v2 in nums2:
                if len(nums)<k:
                    heapq.heappush(nums, [-v1-v2, v1, v2])
                else:
                    if nums[0][0] > -v1-v2:
                        break
                    heapq.heappush(nums, [-v1-v2, v1, v2])
                    heapq.heappop(nums)
        return sorted([[v[1], v[2]] for v in nums])

Solution:

// https://code.dennyzhang.com/find-k-pairs-with-smallest-sums
// Basic Ideas: heap
//
//  Notice: nums1 and nums2 may not have enough candidates
//
// Complexity: Time O(n*m*log(k)), Space O(k)
import "container/heap"
type MyNode struct {
    x, y int
}

// max heap
type IntHeap []MyNode
func (h IntHeap) Len() int { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i].x+h[i].y > h[j].x+h[j].y }
func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x interface{}) {
    *h = append(*h, x.(MyNode))
}
func (h *IntHeap) Pop() interface{} {
    res := (*h)[len(*h)-1]
    *h = (*h)[0:len(*h)-1]
    return res
}

func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
    h := &IntHeap{}
    heap.Init(h)
    for _, x := range nums1 {
        for _, y := range nums2 {
            heap.Push(h, MyNode{x, y})
            if h.Len() > k { heap.Pop(h) }
        }
    }
    size := len(nums1)*len(nums2)
    if size > k { size = k }
    res := make([][]int, size)
    i:=size-1
    for h.Len() != 0 {
        item := heap.Pop(h).(MyNode)
        res[i] = []int{item.x, item.y}
        i--
    }
    return res
}

Post Views: 0