Leet Code: Problem #1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

Problem Statement:

Given an matrix of size m * n which is sorted in rows wise in increasing order. Return the Kth smallest sum of the elements by choosing by atmost 1 element from each row.

Example 1:

Input:

mat = [

            [1, 2, 3],

            [4, 5, 6],

            [7, 8, 9]

            ]

k = 2;

Output:

13 (sum of(2,4,7))

Example 2:

Input:

mat = [

            [1, 2, 3],

            [4, 5, 6],

            [7, 8, 9]

            ]

k = 3;

Output:

14(sum of (3, 4, 7))

Approach to the solution:

1. Take the first row and initialize to an 1D vector
2. From second row on wards add all possible sums of first row and second row and among all those possible sums consider sums that are of length k only(because the sums after that are not worthy considering).
3. Continue the step 2 until all rows are done.
4. The final solution would be the last element of the final 1D vector.

Solution in C++:

class Solution {

public:

    int kthSmallest(vector<vector<int>>& mat, int k) {

        vector<int> sums = mat[0];

        for(int iter = 1; iter < mat.size(); ++iter){

            vector<int> new_sums;

            for(int i = 0; i < mat[iter].size(); ++i){                

                for (int j = 0; j < sums.size(); ++j)    

                    new_sums.push_back(sums[j] + mat[iter][i]);                

            }           

            sort(new_sums.begin(), new_sums.end());            

            new_sums.resize(min(k, (int)new_sums.size()));            

            sums.swap(new_sums);

        }

        return sums.back();

    }

};