Find the Kth smallest element in a matrix of sorted rows-CodePudding

This is an interview question.

Find the Kth smallest element in a matrix of sorted rows, but NOT sorted columns and no relations between the rows. (the first row and nth row have no relation between them - all that is known is that each row is in ascending order)

An example input is this:

[[1,50,60],
 [20,30,40],
 [2,3,4]]
k = 5

And the output in this scenario would be

because 20 is the 5th smallest element in this matrix.

I first thought of adding all the elements to a minHeap, then polling the elements while subtracting one from k each iteration until we have our answer. I also thought about implementing quickselect for a O(m*n) solution, although these solutions dont really take advantage of the fact that all the rows are sorted in ascending order.

What is the optimal way to solve this problem? After thinking about it, I realize it seems like a 'merge k sorted lists' question, where we stop after we find the kth smallest element.

Thanks

CodePudding user response：

Assume the following case

[[1,3,4,...10],
[1,3,4,...10],
[1,2,4,...10]]
For k = 4, ans = 2

Since each row has the same starting/ending element and there's no relation between any two rows, it's hard to determine any useful information without using some brute forcing. The simplest solution should be

Have a max heap with the first k elements of the matrix
Start traversing all rows. If current element is less than max_heap.top, pop the heap and add the current value to it. Else, stop the iteration for the current row (since we can't have the kth smallest element here)

This has minor improvements over complete bruteforcing, but still has O(m*n) time and O(k) space complexity.

CodePudding user response：

Sort the matrix by row, followed by iterating over the row element k times.

Time complexity: n(log(n)   k
Space complexity: n * m
n = number of rows
m = number of columns

public class KthSmallestElementMatrix {

    public static void main(String[] args) throws Exception {

        int[][] matrix = {
          {1, 50, 60},
          {20, 30, 40},
          {2, 3, 4}
        };
        int k = 5;

        System.out.println(kthSmallestElement(matrix, k));
    }

    static class ArrayElement implements Comparable<ArrayElement> {

        int[] arr;

        int index;

        ArrayElement(int[] arr) {
            this.arr = arr;
            this.index = 0;
        }

        public boolean hasNext() {
            return index < arr.length;
        }

        public int next() {

            if (hasNext()) {
                return arr[index  ];
            } else {
                return arr[index];
            }
        }

        @Override
        public int compareTo(ArrayElement o) {
            return Integer.compare(this.arr[index % arr.length], o.arr[o.index % arr.length]);
        }
    }

    private static int kthSmallestElement(int[][] matrix, int k) throws Exception {

        Queue<ArrayElement> q = new PriorityQueue<ArrayElement>();

        for (int i = 0; i < matrix.length; i  ) {
            q.add(new ArrayElement(matrix[i]));
        }

        int element = Integer.MIN_VALUE;

        for (int i = 0; i < k; i  ) {
            ArrayElement ae = q.poll();

            element = ae.next();

            if (ae.hasNext()) {
                q.add(ae);
            }
        }

        return element;
    }
}

CodePudding user response：

One approach that does take advantage of the sorted rows is nested binary search, which is faster when the range of values isn't too large. If your matrix has m rows and n columns, and some Range equal to max(Matrix)-min(Matrix), the time complexity can be written as

O(m * (log n) * (log Range))

The idea behind the algorithm is just the definition of kth smallest element as

"The value such that k-1 elements of the matrix are smaller than the value."

In a matrix, this becomes

"The value such that the sum over each row of the matrix of (the number of elements smaller than value) is equal to k-1".

In a single row, we can count how many elements are smaller than some x using binary search in log n time. For the whole matrix, this takes m * log n time. Initially, you should scan the first and last columns of the matrix to find the smallest and largest values in the matrix, respectively, for your outer binary search bounds.

Here's the pseudocode for the rest of the algorithm:

Scan the first column of matrix and find the minimum; call this low
Scan the last column of matrix and find the maximum; call this high
While low < high:
- Let mid = (low high)/2
- Count the number of elements smaller than mid in the matrix, using binary search
  - If this count is less than k-1: Set low = mid 1
  - Else: Set high = mid
Return low