<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>

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        m, n = len(mat), len(mat[0])
  
        def one_pair(curr, prev):
            heap = []
            heappush(heap, (curr[0] + prev[0], 0, 0))
            count = 0
            res = []
            while heap and count < k:
                val, i, j = heappop(heap)
                res.append(val)
                if j == 0 and i+1 < len(curr):
                    heappush(heap, (curr[i+1] + prev[0], i+1, 0))
                if j+1 < len(prev):
                    heappush(heap, (curr[i] + prev[j+1], i, j+1))
                count += 1
            return res

        prev = mat[m-1][:k]
        for i in range(m-2, -1, -1):
            prev = one_pair(mat[i], prev)
        return prev[-1]
        
# # heapq
# class Solution:
#     def kthSmallest(self, mat: List[List[int]], k: int) -> int:
#         n_row, n_col = len(mat), len(mat[0])
#         smallest_sum = sum(mat[i][0] for i in range(n_row))
#         ids = [0] * n_row
#         heap = [[smallest_sum] + ids]
#         visited = set()

#         for _ in range(k):
#             print(heap)
#             item = heapq.heappop(heap)
#             print(item)
#             for i in range(n_row):
#                 if item[i+1] == n_col-1:
#                     continue
#                 next_item = item.copy()
#                 next_item[0] -= mat[i][next_item[i+1]]
#                 next_item[i+1] += 1
#                 next_item[0] += mat[i][next_item[i+1]]
#                 if tuple(next_item) in visited:
#                     print("check")
#                     continue
#                 visited.add(tuple(next_item))
#                 heapq.heappush(heap, next_item)
#         return item[0]
        
61