LeetCode 1439 - Find the Kth Smallest Sum of a Matrix With Sorted Rows

The problem requires finding the kth smallest sum obtainable by selecting exactly one element from each row of a matrix mat with sorted rows. Each row is sorted in non-decreasing order, and we must explore combinations of elements across rows to form sums.

leetcodehardarraybinary-searchheap-(priority-queue)matrix

LeetCode Problem 1439

Difficulty: 🔴 Hard
Topics: Array, Binary Search, Heap (Priority Queue), Matrix

Solution

Problem Understanding

The problem requires finding the kth smallest sum obtainable by selecting exactly one element from each row of a matrix mat with sorted rows. Each row is sorted in non-decreasing order, and we must explore combinations of elements across rows to form sums. The input is a 2D matrix of dimensions m x n and an integer k. The output is a single integer representing the kth smallest sum of all valid combinations.

Constraints tell us the matrix can have up to 40 rows and 40 columns (1 <= m, n <= 40), each element is positive and at most 5000, and k is bounded by min(200, n^m). This means brute-force generation of all possible combinations is impractical for large inputs because n^m can be enormous, but k is relatively small, allowing optimizations like priority queues or binary search over sums. Important edge cases include matrices with minimal dimensions (1x1), matrices where all elements are the same, and matrices where k is larger than the number of distinct sums.

Approaches

The brute-force approach generates all possible arrays by picking one element from each row, computes their sums, sorts the sums, and returns the kth smallest. While correct, this is inefficient because it requires generating n^m sums, which is computationally infeasible for large matrices.

The key observation for an optimal approach is that rows are sorted, so for any given row, the smallest sum combinations involve taking smaller elements first. This allows us to use a min-heap (priority queue) to efficiently generate the smallest sums in order without enumerating all possibilities. At each step, we expand the smallest sum in the heap by replacing an element in one row with the next element in that row, pushing the new sum into the heap, and maintaining only the k smallest sums. This is a classic application of a k-way merge using a min-heap.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^m log(n^m)) O(n^m) Generate all sums, sort, and pick kth
Min-Heap Merge O(m * k * log k) O(k) Merge rows progressively, keeping only k smallest sums

Algorithm Walkthrough

  1. Initialize a min-heap with the first row of the matrix. This represents the smallest sums possible with only that row.
  2. Iterate over the remaining rows in the matrix. For each row, generate new candidate sums by adding each element of the current row to the sums in the heap from the previous step.
  3. Maintain a min-heap of size at most k for these new sums. This ensures we only keep track of the smallest possible sums.
  4. After processing all rows, the min-heap contains the k smallest sums from all combinations.
  5. Return the largest element in the heap if you maintain a max-heap for the top k sums, or pop from a min-heap k times if using a standard min-heap.

Why it works: By merging rows iteratively and always considering only the next smallest elements due to the sorted property, the algorithm guarantees that the heap will always contain the k smallest sums after each merge. This maintains correctness without enumerating all possible sums.

Python Solution

from typing import List
import heapq

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        # Initialize with the first row
        heap = mat[0][:]  # Make a copy
        heapq.heapify(heap)
        
        for i in range(1, len(mat)):
            new_heap = []
            row = mat[i]
            for prev_sum in heap:
                for num in row:
                    new_sum = prev_sum + num
                    heapq.heappush(new_heap, new_sum)
            # Keep only the k smallest sums
            heap = heapq.nsmallest(k, new_heap)
        
        # The kth smallest sum
        return heap[k - 1]

Implementation Walkthrough: We begin by copying the first row into a heap. For each subsequent row, we compute all possible sums of previous sums plus elements of the current row. To prevent the heap from growing too large, we keep only the k smallest sums at each stage using heapq.nsmallest. After all rows are merged, the kth element in the heap corresponds to the kth smallest sum.

Go Solution

import (
    "container/heap"
    "sort"
)

func kthSmallest(mat [][]int, k int) int {
    heapSums := append([]int(nil), mat[0]...)
    
    for i := 1; i < len(mat); i++ {
        row := mat[i]
        newHeap := []int{}
        for _, prevSum := range heapSums {
            for _, num := range row {
                newHeap = append(newHeap, prevSum+num)
            }
        }
        sort.Ints(newHeap)
        if len(newHeap) > k {
            newHeap = newHeap[:k]
        }
        heapSums = newHeap
    }
    
    return heapSums[k-1]
}

Go-specific notes: Since Go does not have a convenient nsmallest function like Python, we append all new sums into a slice, sort it, and truncate it to the first k elements. This maintains the same correctness invariant as the Python solution. We also make a copy of the first row to avoid modifying the input matrix.

Worked Examples

Example 1: mat = [[1,3,11],[2,4,6]], k = 5

Step by step:

  1. Start with heap [1,3,11].
  2. Merge row [2,4,6]:
  • [1+2, 1+4, 1+6, 3+2, 3+4, 3+6, 11+2, 11+4, 11+6] = [3,5,7,5,7,9,13,15,17]
  • Keep 5 smallest: [3,5,5,7,7]
  1. The 5th smallest sum is 7.

Example 2: mat = [[1,3,11],[2,4,6]], k = 9

  • All sums after merge: [3,5,5,7,7,9,9,11,13,15,17]
  • The 9th smallest sum is 17.

Example 3: mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7

  • Merge row by row maintaining k smallest:
  1. Start [1,10,10]
  2. Merge [1,4,5] → [2,5,6,11,14,15] → keep 7 smallest [2,5,6,11,14,15]
  3. Merge [2,3,6] → [4,5,6,7,8,8,9,10,11,13,14,17,16,19,21,13,16,17] → keep 7 smallest [4,5,5,6,6,7,8]
  • 7th smallest sum: 9.

Complexity Analysis

Measure Complexity Explanation
Time O(m * k * n log k) For each row, we combine k previous sums with n elements, keep only k smallest using heapq.nsmallest or sort+slice
Space O(k) Only k smallest sums are stored at any time

The main savings comes from restricting attention to the k smallest sums instead of all possible sums, reducing exponential growth to manageable linear growth with respect to k.

Test Cases

# Provided examples
assert Solution().kthSmallest([[1,3,11],[2,4,6]], 5) == 7  # Example 1
assert Solution().kthSmallest([[1,3,11],[2,4,6]], 9) == 17  # Example 2
assert Solution().kthSmallest([[1,10,10],[1,4,5],[2,3,6]], 7) == 9  # Example 3

# Edge cases
assert Solution().kthSmallest([[1]], 1) == 1  # Single element
assert Solution().kthSmallest([[1,2],[3,4]], 4) == 6  # k equals total combinations
assert Solution().kthSmallest([[1,2,3],[1,2,3]], 5) == 5  # Multiple equal sums
Test Why
Example 1 Basic correctness and merging of two rows
Example 2 Larger k, tests heap truncation
Example 3 Multiple rows and verifying correct sum order
Single element Minimal input
k equals total combinations Ensures algorithm handles upper bound of k
Multiple equal sums Confirms algorithm handles duplicate sums correctly

Edge Cases

  1. Single-row or single-column matrices: These are trivial, but if mishandled could break the heap logic. The implementation correctly initializes