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

"""
___
Goal:
___

return k-th smallest array sum among possible arrays

___
Examples:
___

Input: mat = [
    [1,3,11],
    [2,4,6]
    ], k = 5

Output: 7

Input: mat = [
    [1,3,11],
    [2,4,6]
    ], k = 9

Output: 17

___
Assumptions:
___

1. we choose exactly one element from each row to form an array

2. rows sorted in increasing order

3. 1 <= mat[i][j] <= 5000

4. 1 <= rows, cols <= 40

___
Algorithms:
___

Algo 1: Compute All Arrays + sort them + index into k-th 
T: O(rows*cols^2 * log(rows*cols^2))
S: O(rows*cols^2)

Algo 2: Compute All Arrays + use k size heap 
T: O(rows*cols^2 * log(rows*k))
S: O(rows*k)


Algo 3: Compute All Arrays + use k size heap <- Possibly use binary search?? 
T: O(rows*cols^2 * log(rows*k))
S: O(rows*k)


Algo 4: Using Heap to Find Next
T: O(k * rows * (cols + log(rows * cols)))
S: O(k * rows * cols)

___
PsuedoCode:
___

[
    [1,10,10],
    [1,4,5],
    [2,3,6]
]

"""


class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:

        ROWS, COLS = len(mat), len(mat[0])
        min_heap: list[tuple[int, list[int]]] = [
            (sum(row[0] for row in mat), [0] * ROWS)
        ]

        seen = set()
        while k - 1:

            curr_sum, indicies = heappop(min_heap)
            k -= 1

            for i in range(ROWS):
                if indicies[i] + 1 < COLS:
                    indicies_cpy = indicies.copy()
                    indicies_cpy[i] += 1
                    next_sum = (
                        curr_sum + mat[i][indicies_cpy[i]] - mat[i][indicies_cpy[i] - 1]
                    )
                    if tuple(indicies_cpy) not in seen:
                        seen.add(tuple(indicies_cpy))
                        heappush(min_heap, (next_sum, indicies_cpy))

        return min_heap[0][0]

121