Leetcode - Find K Pairs with Smallest Sums

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

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 (u₁,v₁),(u₂,v₂) ...(u_k,v_k) 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]

type pair struct {
    idx1 int
    idx2 int
    val int
}
// An IntHeap is a min-heap of ints.
type pairHeap []pair

func (h pairHeap) Len() int           { return len(h) }
func (h pairHeap) Less(i, j int) bool { return h[i].val < h[j].val }
func (h pairHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *pairHeap) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(pair))
}

func (h *pairHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
    if len(nums1) == 0 || len(nums2) == 0 || k == 0 {
        return [][]int{}
    }
    h := &pairHeap{}
    heap.Init(h)
    for i := 0; i < len(nums2); i++ {
        heap.Push(h, pair{0, i, nums1[0]+nums2[i]})
    }
    result := [][]int{}
    if k > len(nums1) * len(nums2) {
        k = len(nums1)*len(nums2)
    }
    for i := 0; i < k; i++ {
        p := heap.Pop(h).(pair)
        result = append(result, []int{nums1[p.idx1], nums2[p.idx2]})
        if p.idx1 < len(nums1)-1{
            heap.Push(h, pair{p.idx1+1, p.idx2, nums1[p.idx1+1]+nums2[p.idx2]})
        }
    }
    return result
}

• 相关文章
• 全局相关文章