<p>Given an <code>n x n</code> <code>matrix</code> where each of the rows and columns is sorted in ascending order, return <em>the</em> <code>k<sup>th</sup></code> <em>smallest element in the matrix</em>.</p>

<p>Note that it is the <code>k<sup>th</sup></code> smallest element <strong>in the sorted order</strong>, not the <code>k<sup>th</sup></code> <strong>distinct</strong> element.</p>

<p>You must find a solution with a memory complexity better than <code>O(n<sup>2</sup>)</code>.</p>

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

<pre>
<strong>Input:</strong> matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
<strong>Output:</strong> 13
<strong>Explanation:</strong> The elements in the matrix are [1,5,9,10,11,12,13,<u><strong>13</strong></u>,15], and the 8<sup>th</sup> smallest number is 13
</pre>

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

<pre>
<strong>Input:</strong> matrix = [[-5]], k = 1
<strong>Output:</strong> -5
</pre>

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

<ul>
	<li><code>n == matrix.length == matrix[i].length</code></li>
	<li><code>1 &lt;= n &lt;= 300</code></li>
	<li><code>-10<sup>9</sup> &lt;= matrix[i][j] &lt;= 10<sup>9</sup></code></li>
	<li>All the rows and columns of <code>matrix</code> are <strong>guaranteed</strong> to be sorted in <strong>non-decreasing order</strong>.</li>
	<li><code>1 &lt;= k &lt;= n<sup>2</sup></code></li>
</ul>

<p>&nbsp;</p>
<p><strong>Follow up:</strong></p>

<ul>
	<li>Could you solve the problem with a constant memory (i.e., <code>O(1)</code> memory complexity)?</li>
	<li>Could you solve the problem in <code>O(n)</code> time complexity? The solution may be too advanced for an interview but you may find reading <a href="http://www.cse.yorku.ca/~andy/pubs/X+Y.pdf" target="_blank">this paper</a> fun.</li>
</ul>

from collections import deque
import math

class Solution:
    def kthSmallest(self, matrix: list[list[int]], k: int) -> int:
        
        n = len(matrix)
        # print(f'n={n} k={k}')
        q = deque()  # queue
        ld = None  # last diagonal added 
        lass = None  # last min assigned
        lminadd = -math.inf  # last min added
        lmaxadd = math.inf  # last max added
        lk = 0  # last 'k' processed

        # Case 1
        if k <= math.ceil(n*n / 2):
    
            while lk < k:
                lk += 1
                ld = (lk-1, 0)
                # print(f'lk={lk} ld={ld}')
                ldmin = math.inf  # diagonal min
                
                # adding a diagonal to queue
                for i in range(0, lk):
                    for j in range(lk-1, -1, -1):
                        if i <= n-1 and j <= n-1 and i+j == lk-1:
                            # print(f'i={i} j={j}')
                            el = matrix[i][j]
                            q.append(el)

                            # keeping track
                            if el < ldmin:
                                ldmin = el

                            # sorting queue
                            prev = 1
                            while 1 + prev <= len(q) and q[-prev] < q[-1-prev]:
                                q[-1-prev], q[-prev] = q[-prev], q[-1-prev]
                                prev += 1
                # print(f'q={q}')
                
                # last min added (from last diagonal)
                lminadd = ldmin
                # print(f'lminadd={lminadd}')

                # assigning values for 'lk'
                lass = q.popleft()
                # print(f'lk={lk} lass={lass}')
                # print(f'q={q}')

                # shortcut to answer
                if 0 <= k - lk - 1 < len(q) and lminadd > lass and lminadd >= q[k - lk - 1]:
                    return q[k - lk - 1]
            
            return lass

        # Case 2
        if k > math.ceil(n*n / 2):

            nk = n*n - k + 1  # new 'k' going backwards
    
            while lk < nk:
                lk += 1
                ld = (n-lk, n-1)
                # print(f'lk={lk} ld={ld}')
                ldmax = -math.inf  # diagonal max
                
                # adding a diagonal to queue
                for i in range(n-lk, n):
                    for j in range(n-1, n-lk-1, -1):
                        if 0 <= i <= n-1 and 0 <= j <= n-1 and i+j == n-lk + n - 1:
                            # print(f'i={i} j={j}')
                            el = matrix[i][j]
                            q.append(el)

                            # keeping track
                            if el > ldmax:
                                ldmax = el

                            # sorting queue
                            prev = 1
                            while 1 + prev <= len(q) and q[-prev] > q[-1-prev]:
                                q[-1-prev], q[-prev] = q[-prev], q[-1-prev]
                                prev += 1
                # print(f'q={q}')
                
                # last max added (from last diagonal)
                lmaxadd = ldmax
                # print(f'lmaxadd={lmaxadd}')

                # assigning values for 'lk'
                lass = q.popleft()
                # print(f'lk={lk} nk={n*n-lk+1} lass={lass}')
                # print(f'q={q}')

                # jumping to answer
                if 0 <= nk - lk - 1 < len(q) and lmaxadd < lass and lmaxadd <= q[nk - lk - 1]:
                    return q[nk - lk - 1]

            return lass


# Working simple non-efficient solution
# class Solution:
#     def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
#         n = len(matrix[0])
#         arr = []
#         for i in range(n):
#             for j in range(n):
#                 arr.append(matrix[i][j])
#         arr.sort()
#         return arr[k-1]
    
21400