Scale Engineer

Find the Kth Smallest Sum of a Matrix With Sorted Rows

Hard
#1331Time: O(n^m * log(n^m)). It takes O(n^m) to generate all the sums, and then O(n^m * log(n^m)) to sort them. The sorting step dominates.Space: O(n^m) to store all the generated sums in a list before sorting.
Algorithms
Binary Search
Data structures
ArrayHeap (Priority Queue)Matrix

Prompt

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

You are allowed to choose exactly one 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.

 

Constraints:

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

Approaches

4 approaches with complexity analysis and trade-offs.

The most straightforward approach is to generate every possible array sum. There are m rows and n choices in each row, leading to a total of n^m possible arrays. We can use a recursive backtracking function to explore all these combinations. As we form each complete array, we calculate its sum and store it in a list. After generating all possible sums, we sort the list and pick the k-th element (at index k-1).

Algorithm

  • Define a recursive function, say generateSums(row, currentSum), that explores all possible combinations.
  • The function takes the current row index and the sum accumulated so far.
  • Base Case: If the row index reaches m (the number of rows), it means we have picked one element from each row. Add the currentSum to a global list of all possible sums.
  • Recursive Step: For the current row, iterate through each element mat[row][j]. Make a recursive call generateSums(row + 1, currentSum + mat[row][j]) for each element.
  • Start the process by calling generateSums(0, 0).
  • After the recursion completes, the list will contain all n^m possible sums.
  • Sort this list in non-decreasing order.
  • The k-th smallest sum is the element at index k-1 in the sorted list.

Walkthrough

This method exhaustively enumerates all possibilities. A recursive function can traverse the matrix row by row. At each row, it branches out for every element in that row, adding it to the sum being built. When it reaches the end of the matrix (after picking one element from the last row), the accumulated sum is one of the possible array sums and is stored. Once all n^m sums are collected, sorting them allows us to find the k-th smallest one.

import java.util.ArrayList;import java.util.Collections;import java.util.List; class Solution {    List<Integer> allSums;    int m, n;    int[][] mat;     public int kthSmallest(int[][] mat, int k) {        this.allSums = new ArrayList<>();        this.mat = mat;        this.m = mat.length;        this.n = mat[0].length;         generateSums(0, 0);                Collections.sort(allSums);                return allSums.get(k - 1);    }     private void generateSums(int row, int currentSum) {        if (row == m) {            allSums.add(currentSum);            return;        }         for (int j = 0; j < n; j++) {            generateSums(row + 1, currentSum + mat[row][j]);        }    }}

Complexity

Time

O(n^m * log(n^m)). It takes O(n^m) to generate all the sums, and then O(n^m * log(n^m)) to sort them. The sorting step dominates.

Space

O(n^m) to store all the generated sums in a list before sorting.

Trade-offs

Pros

  • Simple to conceptualize and implement.

Cons

  • Extremely high time complexity, making it infeasible for the given constraints.

  • Requires a large amount of memory to store all possible sums, which can lead to memory limit errors.

Solutions

class Solution {public  int kthSmallest(int[][] mat, int k) {    int m = mat.length, n = mat[0].length;    List<Integer> pre = new ArrayList<>(k);    List<Integer> cur = new ArrayList<>(n * k);    pre.add(0);    for (int[] row : mat) {      cur.clear();      for (int a : pre) {        for (int b : row) {          cur.add(a + b);        }      }      Collections.sort(cur);      pre.clear();      for (int i = 0; i < Math.min(k, cur.size()); ++i) {        pre.add(cur.get(i));      }    }    return pre.get(k - 1);  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem