Kth Smallest Element in a Sorted Matrix in Python

Finding the Kth smallest element in a sorted matrix is a classic problem that can be solved efficiently using binary search. In this problem, we have an n x n matrix where each row and column is sorted in increasing order, and we need to find the kth smallest element in the entire matrix.

Problem Understanding

Given a matrix where rows and columns are sorted in ascending order, we need to find the kth smallest element. For example, in the matrix [[1,5,9],[10,11,13],[12,13,15]], if k=8, the answer is 13 because when all elements are sorted: [1,5,9,10,11,12,13,13,15], the 8th element is 13.

Algorithm Approach

We use binary search on the value range rather than indices. The key insight is to count how many elements are smaller than or equal to a given value using the sorted property of the matrix ?

Step-by-Step Process

• Set low to the top-left element and high to the bottom-right element
• For each mid value, count elements ≤ mid using the check_value function
• If count < k, search in the higher range; otherwise, search in the lower range
• The check_value function starts from the top-right corner and moves left or down based on comparison

Implementation

class Solution:
    def kthSmallest(self, matrix, k):
        """
        Find the kth smallest element in a sorted matrix
        :param matrix: List[List[int]] - sorted matrix
        :param k: int - position of smallest element to find
        :return: int - kth smallest element
        """
        n = len(matrix)
        high = matrix[n-1][n-1]  # Bottom-right element
        low = matrix[0][0]       # Top-left element
        
        while low <= high:
            mid = low + (high - low) // 2
            count = self.check_value(matrix, mid)
            
            if count < k:
                low = mid + 1
            else:
                high = mid - 1
        
        return low
    
    def check_value(self, matrix, value):
        """
        Count elements smaller than or equal to value
        """
        i = 0  # Start from top row
        j = len(matrix[0]) - 1  # Start from rightmost column
        counter = 0
        
        while i < len(matrix) and j >= 0:
            if matrix[i][j] > value:
                j -= 1  # Move left
            else:
                counter += j + 1  # Count all elements in current row up to j
                i += 1  # Move down
        
        return counter

# Example usage
matrix = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
solution = Solution()
result = solution.kthSmallest(matrix, 8)
print(f"The 8th smallest element is: {result}")

The 8th smallest element is: 13

How the Algorithm Works

The check_value function efficiently counts elements by starting from the top-right corner. If the current element is greater than our target value, we move left. If it's smaller or equal, we count all elements in that row up to the current position and move down.

Example Walkthrough

# Let's trace through the algorithm step by step
matrix = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
k = 8

solution = Solution()

# Initial values
low = matrix[0][0]    # 1
high = matrix[2][2]   # 15

print(f"Initial range: low={low}, high={high}")

# First iteration: mid = 8
mid = (1 + 15) // 2  # 8
count = solution.check_value(matrix, 8)
print(f"Mid={mid}, Count of elements <= {mid}: {count}")

# Since count=3 < k=8, we search higher
low = mid + 1  # 9
print(f"New range: low={low}, high={high}")

Initial range: low=1, high=15
Mid=8, Count of elements <= 8: 3
New range: low=9, high=15

Time and Space Complexity

• Time Complexity: O(n × log(max-min)) where n is the matrix dimension
• Space Complexity: O(1) as we use only constant extra space

Conclusion

This binary search approach efficiently finds the kth smallest element by searching on values rather than indices. The key insight is using the sorted property to count elements smaller than a given value in O(n) time, making the overall solution optimal for large matrices.