Kth Smallest Element in a Sorted Matrix

This is Thursday and a template summer day in the Twin Cities of Minneapolis and St. Paul. The highs for today have been forecasted in the low to mid 70s F. Overall the temperatures a lower than usual.

Earlier this morning I watched the NewScientist DEBATE “Information and the future of defense” presented by BAE Systems. It was interesting. The panel had a moderator and if I recall four contributors with very good credentials in the topic at hand.

In my humble opinion, cyberwarfare is a very serious threat brought by bad actors which may represent organizations or states. Since they can be very damaging and are constantly morphing, AI should be employed to determine the threat in a very short time (e.g., milliseconds), block it, and bring harm to the actors. There is an old saying, if force does not solve your problem, you are not using enough. I will not go into more detail at this time, but chatting, discussing and signing agreements between good actors will never stop bad ones.

I decided to work on LeetCode 378 Kth Smallest Element in a Sorted Matrix problem. The problem is classified as Medium. It is quite interesting because after reading and assimilating the requirements, it looked like the solution could be a modified binary search. That is easier said than done.

Given an n x n matrix where each of the rows and columns are 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.

Constraints:

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

We are given a square matrix of n * n and a number `k`. We need to return the kth smallest element in the matrix. If you need a few minutes to put your thoughts together, this is a good time to do so. I will wait until you are ready…

…OK now that you are back let’s take a look at the signature of the function of interest. We are going to solve the problem using the Java programming language and the VSCode IDE on a Windows computer.

    public int kthSmallest(int[][] matrix, int k) {
        
    }

We are provided the sorted matrix and the value of the kth element we need to locate in the matrix.

8
1,5,9
10,11,13
12,13,15
main <<< k: 8
main <<< n: 3
main <<< matrix:
[1, 5, 9]
[10, 11, 13]
[12, 13, 15]
main <<< kthSmallest0: 13
main <<<  kthSmallest: 13


1
-5
main <<< k: 1
main <<< n: 1
main <<< matrix:
[-5]
main <<< kthSmallest0: -5
main <<<  kthSmallest: -5


3
10,20,30,40
15,25,35,45
24,29,37,48
32,33,39,50
main <<< k: 3
main <<< n: 4
main <<< matrix:
[10, 20, 30, 40]
[15, 25, 35, 45]
[24, 29, 37, 48]
[32, 33, 39, 50]
main <<< kthSmallest0: 20
main <<<  kthSmallest: 20

The first input line contains the value for `k`. The following few lines contain the contents of the matrix. Note that we are dealing with square matrices.

Our test code, which IS NOT PART OF THE SOLUTION, reads and display the value of `k`. It then figures the value of `n` and displays it on the screen.

The program then reads and displays the contents of the int[][] `matrix`.

At this point it seems that two implementations of the function of interest are called. Both seem to return the same results.

    /**
     * Test scaffold
     * @throws IOException
     */
    public static void main(String[] args) throws IOException {
        
        // **** open buffered reader ****
        BufferedReader br = new BufferedReader(new java.io.InputStreamReader(System.in));

        // **** read k ****
        int k = Integer.parseInt(br.readLine().trim());

        // **** read first line for matrix ****
        String[] line = br.readLine().trim().split(",");

        // **** extract value for n ****
        int n = line.length;

        // **** declare matrix ****
        int[][] matrix = new int[n][n];
        matrix[0] = Arrays.stream(line).mapToInt(Integer::parseInt).toArray();

        // **** populate the rest of the matrix ****
        for (int i = 1; i < n; i++) {
            matrix[i] = Arrays.stream(br.readLine().trim().split(","))
                            .mapToInt(Integer::parseInt)
                            .toArray();
        }
    
        // **** close buffered reader ****
        br.close();

        // ???? ????
        System.out.println("main <<< k: " + k);
        System.out.println("main <<< n: " + n);
        System.out.println("main <<< matrix:");
        for (int i = 0; i < n; i++)
            System.out.println(Arrays.toString(matrix[i]));

        // ???? ????
        // for (int i = 0; i < 17; i++) {

        //     // ???? ????
        //     Random rand = new Random();

        //     // ???? generate mid ????
        //     int mid = matrix[0][0] + rand.nextInt(matrix[n - 1][n - 1] - matrix[0][0] + 1);
        //     if (mid < matrix[0][0] || mid > matrix[n - 1][n - 1]) {
        //         System.err.println("main <<< invalid mid: " + mid); System.exit(-1);
        //     }

        //     // ???? compute count of mid ????
        //     System.out.println("main <<< count(" + mid + "): " + count0(matrix, mid));
        // }

        // **** call function of interest and display result ****
        System.out.println("main <<< kthSmallest0: " + kthSmallest0(matrix, k));

        // **** call function of interest and display result ****
        System.out.println("main <<<  kthSmallest: " + kthSmallest(matrix, k));
    }

The code for the test scaffold seems to follow the description to process a test case.

Note that there is some code commented out. This was done to test the count0() function. We will learn more about the actual function when we go over its associated code.

    /**
     * Given an n x n matrix where each of the rows and columns are 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.
     * 
     * 85 / 85 test cases passed.
     * Status: Accepted
     * Runtime: 0 ms
     * Memory Usage: 44.7 MB
     */
    static int kthSmallest0 (int[][] matrix, int k) {

        // **** sanity check(s) ****
        if (k == 1) return matrix[0][0];

        // **** initialization ****
        int n       = matrix.length;            // for ease of use
        int low     = matrix[0][0];
        int high    = matrix[n - 1][n - 1];

        // **** binary search ****
        while (low < high) {

            // **** compute mid value ****
            int mid = (high - low) / 2 + low;

            // **** update high or low value ****
            if (count0(matrix, mid) >= k)
                high = mid;
            else 
                low = mid + 1;
        }

        // **** return low ****
        return low;
    }

This is the first implementation of the function of interest.

We start by performing a sanity check. This particular test checks for the case when ‘k’ is equal to one. It returns the lowest value in the `matrix` which should be at coordinates [0][0].

We then perform some initializations.

The while loop implements the binary search.

We first compute the mid value. Once we have it, we update the low or high value.

When all is said and done, we return the low value.

Hopefully you have noted the close resemblance of this code to a traditional binary search. If in doubt you can refer to the Binary Search in Java post in this blog.

    /**
     * Count the # of values that is less than or equal to mid.
     * 
     * Time: O(n ^ 2) - Space: O(1)
     */
    static private int count0 (int[][] matrix, int mid) {

        // **** initialization ****
        int count   = 0;
        int n       = matrix.length;        // for ease of use
        int col     = n - 1;

        // **** traverse rows (top to bottom) - O(n)****
        for (int row = 0; row < n; row++) {

            // **** traverse columns (right to left) - O(n) ****
            for ( ; col >= 0; col--) {

                // **** increment count (if needed) ****
                if (matrix[row][col] <= mid) {

                    // **** increment count ****
                    count += (col + 1);

                    // **** exit inner loop ****
                    break;
                }
            }
        }

        // **** return count ****
        return count;
    }

As we noted in the function of interest, we make a call to the count0() function to determine the way we proceed with the binary search.

We start by initializing a set of variables.

In the first loop we search the rows from top to bottom.

In the inner loop, we search the columns for the specified row right to left. This is done because each row is ordered in ascending order.

If the current value in the `matrix` is <= `mid` we increment the count of values we have so far and exit the inner loop.

When all is said and done, we return the count.

    /**
     * Given an n x n matrix where each of the rows and columns are 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.
     * 
     * 85 / 85 test cases passed.
     * Status: Accepted
     * Runtime: 0 ms
     * Memory Usage: 44.7 MB
     */
    static int kthSmallest (int[][] matrix, int k) {

        // **** sanity check(s) ****
        if (k == 1) return matrix[0][0];
        
        // **** initialization ****
        int n       = matrix.length;        // for ease of use
        int low     = matrix[0][0];
        int high    = matrix[n - 1][n - 1];

        // **** perform binary search ****
        while (low <= high) {

            // **** compute mid ****
            int mid = low + (high - low) / 2;

            // **** update high or low ****
            if (check(matrix, mid, k, n))
                high = mid - 1;
            else
                low = mid + 1;
        }

        // **** return low ****
        return low;
    }

This is the second implementation of the function of interest. It is quite similar to the first one. We start by performing a sanity check and then the initialization step. We then enter the loop which performs the binary search. We compute the value for the `mid` point. Based on the value returned by the check() function we update the low or the high value. When all is said and done we return the low value. This is quite similar to the previous implementation.

    /**
     * Check to update high (return true) or low (return false) value. 
     */
    private static boolean check (int[][] matrix, int mid, int k, int n) {

        // **** initialization ****
        int col = n - 1;
        int row = 0;
        int num = 0;

        // **** move in the matrix updating the num ****
        while (col >= 0 && row < n) {
            if (matrix[col][row] <= mid) {
                num += (col + 1);
                row++;
            } else
                col--;
        }

        // **** return true (updte high) or false (update low) ****
        return num >= k;
    }

This function is used to determine in which direction we move after computing the `mid` value.

We start by performing the initialization step.

In the while loop, we move inside the `matrix` until the `num` value can be updated. At that time we return the Boolean value for the `num >= k` relation. This indicated the direction in which the binary search will continue.

Note that both functions have passed all the tests and the times and space are comparable.

Hope you enjoyed solving this problem as much as I did. The entire code for this project can be found in my GitHub repository.

If you have comments or questions regarding this, or any other post in this blog, please do not hesitate and leave me a note below. I will reply as soon as possible.

Keep on reading and experimenting. It is one of the best ways to learn, become proficient, refresh your knowledge and enhance your developer toolset.

Thanks for reading this post, feel free to connect with me John Canessa at LinkedIn.

Regards;

John

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.