We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Find the Kth Smallest Sum of a Matrix With Sorted Rows

Number: 1550

Difficulty: Hard

Paid? No

Companies: Amazon, Meta


Problem Description

Given an m x n matrix where each row is sorted in non-decreasing order, you are to pick exactly one element from each row to form an array. The goal is to determine the kth smallest sum that can be generated by summing the selected elements across all rows.


Key Insights

  • Each row is already sorted which can be exploited in merging sums.
  • Instead of generating all possible combinations (which is exponential), iteratively merge the rows while only keeping the k smallest sums.
  • For each new row, combine each existing sum with every element of the row and then sort and trim the list to at most k candidates.
  • Using a heap (priority queue) can optimize the merging process by efficiently extracting the smallest sums.
  • This approach leverages the constraint on k to ensure the solution remains efficient even though the matrix dimensions could be relatively large.

Space and Time Complexity

Time Complexity: O(m * k * n * log(k * n)) – For each of m rows, we merge up to k sums with n candidates each, and sorting or heap operations cost logarithmic time. Space Complexity: O(k * n) – At each merge step, we maintain at most k * n intermediate sums before trimming to k values.


Solution

We use an iterative merging strategy:

  1. Start with a list containing a single element: zero, representing the sum before processing any rows.
  2. For each row in the matrix, generate new candidate sums by adding each element of the row to each of the current sums.
  3. Sort the generated list of candidate sums and retain only the k smallest ones, as larger sums are not needed.
  4. After processing all rows, the kth smallest sum is the kth element of this merged list. This method leverages the pre-sorted property of rows and restricts the number of candidates maintained, making it efficient for the given constraints.

Code Solutions

# Python Solution

def kthSmallest(mat, k):
    # Start with an initial list of sums, which only has 0.
    current_sums = [0]
    
    # Iterate over each row in the matrix.
    for row in mat:
        new_sums = []
        # Combine each current sum with each number in the current row.
        for prev_sum in current_sums:
            for num in row:
                new_sums.append(prev_sum + num)
        # Sort and keep only the first k smallest sums.
        new_sums.sort()
        current_sums = new_sums[:k]
    
    # The kth smallest sum is at index k-1 in the final list.
    return current_sums[k - 1]

# Example usage:
# print(kthSmallest([[1,3,11],[2,4,6]], 5))  # Output: 7
← Back to All Questions