Find Kth Smallest Sum of a Matrix with Sorted Rows in C++

Suppose we have one m * n matrix called mat, and an integer k, mat has its rows sorted in nondecreasing order. We can choose exactly one element from each row to form an array. We have to find the Kth smallest array sum among all possible arrays.

So, if the input is like mat = [[1,3,11],[2,4,6]]

and k = 5, then the output will be 7, as when we choose one element from each row the first k smallest sums are [1,2], [1,4], [3,2], [3,4], [1,6]. here the 5th sum is 7.

To solve this, we will follow these steps −

• Define one priority queue pq

• Define one 2D array m

• Define a function update(), this will take an array v, i, ok initialize it with false,

• if i is same as size of v, then −

  • if ok is false, then −

    • return

  • return

  • for initialize j := 0, when j < size of v, update (increase j by 1), do −

    • sum := sum + m[j, v[j]]

  • Define an array temp and copy v into it

  • insert sum into temp at the beginning

  • insert temp into pq

  • return

• (increase v[i] by 1)

• if ok is false and v[i] < z, then −

  • update(v, i + 1, true)

• update(v, i + 1, true)

• update(v, i + 1, ok)

• From the main method do the following −

• m :+ given matrix

• ret := 0

• n := row count of m

• z := column count of m

• for initialize i := 0, when i < n, update (increase i by 1), do −

  • ret := ret + m[i, 0]

• Define an array temp of size n

• insert ret into temp at the beginning

• insert temp into pq

• Define one set s

• while k is non-zero, decrease k by 1 in each iteration, do −

  • Define an array temp = top of pq

  • delete element from pq

  • insert temp into s

  • ret := temp[0]

  • delete first element of temp from temp

  • update(temp, 0)

  • while (not pq is empty and top element of pq is member of s), do −

    • delete element from pq

• return ret

Let us see the following implementation to get better understanding −

Example

 Live Demo

#include <bits/stdc++.h>
using namespace std;
struct Cmp{
   bool operator()(vector <int>& a, vector <int>& b) {
      return !(a[0] < b[0]);
   }
};
class Solution {
   public:
   priority_queue<vector<int>, vector<vector<int> >, Cmp> pq;
   vector<vector<int> > m;
   int z;
   void update(vector<int>& v, int i, bool ok = false){
      if (i == v.size()) {
         if (!ok)
         return;
         int sum = 0;
         for (int j = 0; j < v.size(); j++) {
            sum += m[j][v[j]];
         }
         vector<int> temp(v.begin(), v.end());
         temp.insert(temp.begin(), sum);
         pq.push(temp);
         return;
      }
      v[i]++;
      if (!ok && v[i] < z)
      update(v, i + 1, true);
      v[i]--;
      update(v, i + 1, ok);
   }
   int kthSmallest(vector<vector<int> >& m, int k){
      this->m = m;
      int ret = 0;
      int n = m.size();
      z = m[0].size();
      for (int i = 0; i < n; i++) {
         ret += m[i][0];
      }
      vector<int> temp(n);
      temp.insert(temp.begin(), ret);
      pq.push(temp);
      set<vector<int> > s;
      while (k--) {
         vector<int> temp = pq.top();
         pq.pop();
         s.insert(temp);
         ret = temp[0];
         temp.erase(temp.begin());
         update(temp, 0);
         while (!pq.empty() && s.count(pq.top())) {
            pq.pop();
         }
      }
      return ret;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{1,3,11},{2,4,6}};
   cout << (ob.kthSmallest(v, 5));
}

Input

{{1,3,11},{2,4,6}}

Output

7