Kth Smallest in a Matrix

The Idea

Rows and columns are sorted, but the matrix is not one sorted list, so there is no index to binary search. Binary search the values instead. Guess a number between the two corners and count how many entries are at most that guess — the staircase walk does it in O(n). Too few, and the answer is higher; enough, and the answer is this value or lower. The range shrinks to one number, and that number is in the matrix.

Real-World Example

Finding the median pay across sorted salary bands held by different offices, without merging their lists. Name a figure, ask each office how many of its people are below it, and move the figure until the counts line up.

The Code

def kth_smallest(matrix, k):
    n = len(matrix)
    lo, hi = matrix[0][0], matrix[n - 1][n - 1]
    while lo < hi:
        mid = (lo + hi) // 2
        count, c = 0, n - 1
        for r in range(n):                    # staircase count of values <= mid
            while c >= 0 and matrix[r][c] > mid:
                c -= 1
            count += c + 1
        if count < k:
            lo = mid + 1
        else:
            hi = mid
    return lo

print(kth_smallest([[1, 5, 9], [10, 11, 13], [12, 13, 15]], 8))   # 13

M = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
count, c = 0, 2
for r in range(3):
  while c >= 0 and M[r][c] > 11:
      c -= 1
  count += c + 1
print(count)