{x}
blog image

Find the Kth Smallest Sum of a Matrix With Sorted Rows

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.

Solution Explanation

This problem asks to find the kth smallest sum of an array formed by picking exactly one element from each row of a matrix where each row is sorted in non-decreasing order. The naive approach of generating all possible sums and sorting them is computationally expensive, especially with larger matrices. A more efficient solution leverages the sorted nature of the rows and uses a heap or a similar approach to maintain a sorted list of sums.

The solutions presented utilize a dynamic programming approach with a clever optimization. Instead of explicitly generating all possible sums, the algorithm iteratively builds a sorted list of partial sums.

Algorithm:

  1. Initialization: Start with a list pre containing only the element 0. This represents the sum of choosing no elements yet.

  2. Iteration: Iterate through each row of the matrix. For each row:

    • Create a new list cur.
    • For each element a in pre (previous partial sums) and each element b in the current row, calculate the sum a + b and add it to cur. This generates all possible sums when considering the current row in addition to previous rows.
    • Sort cur. This ensures the list remains sorted.
    • Take the first k elements from cur and store them back in pre. This is the crucial optimization: we only keep track of the k smallest sums at any given step, discarding the larger ones as they can't contribute to the kth smallest overall sum.
  3. Result: After processing all rows, the kth smallest element in pre is the solution.

Time Complexity Analysis:

  • The outer loop iterates m times (number of rows).
  • In each iteration, we perform operations on a list with size at most k. The sorting operation dominates, taking O(k log k) time.
  • The generation of cur involves nested loops, but the nested loop iterations are bounded by O(k*n), where n is the number of columns and k is the target k-th smallest sum.
  • Therefore, the overall time complexity is O(m * k * (k + n log k)), approximately O(m * k * n log k). This is significantly better than the O(nm log(nm)) complexity of a brute-force approach.

Space Complexity Analysis:

  • The algorithm uses pre and cur lists. The maximum size of both is O(k), so the space complexity is O(k).

Code Examples (Python):

The Python code directly implements the algorithm described above. It uses the built-in sorted() function for sorting the list cur and list slicing to efficiently get the top k elements.

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        pre = [0]
        for cur in mat:
            pre = sorted(a + b for a in pre for b in cur[:k])[:k] # Optimize by considering only up to k elements of each row
        return pre[-1]
 

The other language examples follow the same approach, adapting the specific syntax and data structures to each language. Note that efficient sorting is crucial for performance. In Java, Collections.sort() is used. C++ uses std::sort(), Go uses sort.Ints(), and Javascript uses sort(). The TypeScript solution utilizes the sort() method for array sorting.

The provided solutions efficiently solve the problem by intelligently limiting the number of sums generated and maintained, resulting in a substantially improved time complexity compared to a brute-force solution.