1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows #

Problem #

You are given an m * n matrix, mat, and an integer k, which has its rows sorted in non-decreasing order.

You are allowed to choose exactly 1 element from each row to form an array. Return the Kth smallest array sum among all possible arrays.

Example 1:

Input: mat = [[1,3,11],[2,4,6]], k = 5
Output: 7
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,2], [1,4], [3,2], [3,4], [1,6]. Where the 5th sum is 7.  

Example 2:

Input: mat = [[1,3,11],[2,4,6]], k = 9
Output: 17

Example 3:

Input: mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
Output: 9
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]. Where the 7th sum is 9.  

Example 4:

Input: mat = [[1,1,10],[2,2,9]], k = 7
Output: 12

Constraints:

• m == mat.length
• n == mat.length[i]
• 1 <= m, n <= 40
• 1 <= k <= min(200, n ^ m)
• 1 <= mat[i][j] <= 5000
• mat[i] is a non decreasing array.

Problem Summary #

Given an m * n matrix mat and an integer k, each row in the matrix is sorted in non-decreasing order. You can choose 1 element from each row to form an array. Return the kth smallest array sum among all possible arrays.

Solution Approach #

• This problem is an upgraded version of Problem 373. In Problem 373, given 2 sorted arrays, we need to choose one number from each of the 2 arrays to form a pair, and finally output the K pairs with the smallest sums. In this problem, we are given an m*n matrix. Essentially, the 2 arrays in Problem 373 are upgraded to m arrays. It is simply one more outer loop. This loop chooses one number from each row in order: first take numbers from row 0 and row 1, find the first K smallest combinations, then take numbers from row 2, and so on. The other methods are consistent with Problem 373. Maintain a min-heap of length k. Each time, pop the smallest array sum sum and the corresponding index index from the heap, then move the index backward by one position in order, generate a new sum, and add it to the heap.

Code #

package leetcode

import "container/heap"

func kthSmallest(mat [][]int, k int) int {
	if len(mat) == 0 || len(mat[0]) == 0 || k == 0 {
		return 0
	}
	prev := mat[0]
	for i := 1; i < len(mat); i++ {
		prev = kSmallestPairs(prev, mat[i], k)
	}
	if k < len(prev) {
		return -1
	}
	return prev[k-1]
}

func kSmallestPairs(nums1 []int, nums2 []int, k int) []int {
	res := []int{}
	if len(nums2) == 0 {
		return res
	}
	pq := newPriorityQueue()
	for i := 0; i < len(nums1) && i < k; i++ {
		heap.Push(pq, &pddata{
			n1:    nums1[i],
			n2:    nums2[0],
			n2Idx: 0,
		})
	}
	for pq.Len() > 0 {
		i := heap.Pop(pq)
		data := i.(*pddata)
		res = append(res, data.n1+data.n2)
		k--
		if k <= 0 {
			break
		}
		idx := data.n2Idx
		idx++
		if idx >= len(nums2) {
			continue
		}
		heap.Push(pq, &pddata{
			n1:    data.n1,
			n2:    nums2[idx],
			n2Idx: idx,
		})
	}
	return res
}

type pddata struct {
	n1    int
	n2    int
	n2Idx int
}

type priorityQueue []*pddata

func newPriorityQueue() *priorityQueue {
	pq := priorityQueue([]*pddata{})
	heap.Init(&pq)
	return &pq
}

func (pq priorityQueue) Len() int           { return len(pq) }
func (pq priorityQueue) Swap(i, j int)      { pq[i], pq[j] = pq[j], pq[i] }
func (pq priorityQueue) Less(i, j int) bool { return pq[i].n1+pq[i].n2 < pq[j].n1+pq[j].n2 }
func (pq *priorityQueue) Pop() interface{} {
	old := *pq
	val := old[len(old)-1]
	old[len(old)-1] = nil
	*pq = old[0 : len(old)-1]
	return val
}

func (pq *priorityQueue) Push(i interface{}) {
	val := i.(*pddata)
	*pq = append(*pq, val)
}