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

<p>You are allowed to choose <strong>exactly one element</strong> from each row to form an array.</p>

<p>Return <em>the </em><code>k<sup>th</sup></code><em> smallest array sum among all possible arrays</em>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> mat = [[1,3,11],[2,4,6]], k = 5
<strong>Output:</strong> 7
<strong>Explanation:</strong> 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.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> mat = [[1,3,11],[2,4,6]], k = 9
<strong>Output:</strong> 17
</pre>

<p><strong class="example">Example 3:</strong></p>

<pre>
<strong>Input:</strong> mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
<strong>Output:</strong> 9
<strong>Explanation:</strong> 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.  
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>m == mat.length</code></li>
	<li><code>n == mat.length[i]</code></li>
	<li><code>1 &lt;= m, n &lt;= 40</code></li>
	<li><code>1 &lt;= mat[i][j] &lt;= 5000</code></li>
	<li><code>1 &lt;= k &lt;= min(200, n<sup>m</sup>)</code></li>
	<li><code>mat[i]</code> is a non-decreasing array.</li>
</ul>

struct HeapElement{
public:
    std::vector<int> colIndices;
    int sum;
    bool largerColIndices(const HeapElement& other) const {
        for (int i = 0; i < this->colIndices.size(); ++i)
            if (this->colIndices[i] < other.colIndices[i])
                return false;
            else if (this->colIndices[i] > other.colIndices[i])
                return true;
        return false;
    }
    bool smallerColIndices(const HeapElement& other) const {
        for (int i = 0; i < this->colIndices.size(); ++i)
            if (this->colIndices[i] > other.colIndices[i])
                return false;
            else if (this->colIndices[i] < other.colIndices[i])
                return true;
        return false;
    }
    bool equalColIndices(const HeapElement& other) const {
        for (int i = 0; i < this->colIndices.size(); ++i)
            if (this->colIndices[i] != other.colIndices[i])
                return false;
        return true;
    }
};

bool operator<(const HeapElement& v1, const HeapElement& v2) {
    return (v1.sum < v2.sum) || (v1.sum == v2.sum && v1.smallerColIndices(v2));
}

bool operator>(const HeapElement& v1, const HeapElement& v2) {
    return (v1.sum > v2.sum) || (v1.sum == v2.sum && v1.largerColIndices(v2));
}


class Solution {
public:
    int kthSmallest(const std::vector<std::vector<int>>& mat, int k) {
        int m = mat.size();
        int n = mat[0].size();
        std::priority_queue<HeapElement, std::vector<HeapElement>, std::greater<HeapElement>> pq{};
        int smallestSum = 0;
        for (int r = 0; r < m; ++r)
            smallestSum += mat[r][0];
        pq.push(HeapElement{
            .colIndices = std::vector<int>(m, 0),
            .sum = smallestSum
        });
        for (int i = 1; i <= k - 1; ++i) {
            HeapElement top = pq.top();
            pq.pop();
            while (!pq.empty() && pq.top().equalColIndices(top))
                pq.pop();
            for (int r = 0; r < m; ++r) {
                if (top.colIndices[r] < n - 1) {
                    HeapElement toBePushed = top;
                    toBePushed.sum += mat[r][toBePushed.colIndices[r] + 1] - mat[r][toBePushed.colIndices[r]];
                    ++toBePushed.colIndices[r];
                    pq.push(toBePushed);
                }
            }
        }
        return pq.top().sum;
    }
};
44191