LeetCode in Kotlin

378. Kth Smallest Element in a Sorted Matrix

Medium

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

You must find a solution with a memory complexity better than O(n2).

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8

Output: 13

Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8^th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1

Output: -5

Constraints:

• n == matrix.length == matrix[i].length
• 1 <= n <= 300
• -109 <= matrix[i][j] <= 109
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= n2

Follow up:

• Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?
• Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

Solution

class Solution {
    fun kthSmallest(matrix: Array<IntArray>?, k: Int): Int {
        if (matrix == null || matrix.size == 0) {
            return -1
        }
        var start = matrix[0][0]
        var end = matrix[matrix.size - 1][matrix[0].size - 1]
        // O(log(max-min)) time
        while (start + 1 < end) {
            val mid = start + (end - start) / 2
            if (countLessEqual(matrix, mid) < k) {
                // look towards end
                start = mid
            } else {
                // look towards start
                end = mid
            }
        }

        // leave only with start and end, one of them must be the answer
        // try to see if start fits the criteria first
        return if (countLessEqual(matrix, start) >= k) {
            start
        } else {
            end
        }
    }

    // countLessEqual
    // O(n) Time
    private fun countLessEqual(matrix: Array<IntArray>, target: Int): Int {
        // binary elimination from top right
        var row = 0
        var col: Int = matrix[0].size - 1
        var count = 0
        while (row < matrix.size && col >= 0) {
            if (matrix[row][col] <= target) {
                // get the count in current row
                count += col + 1
                row++
            } else if (matrix[row][col] > target) {
                // eliminate the current col
                col--
            }
        }
        return count
    }
}

This site is open source. Improve this page.