Find the Kth smallest element in a matrix of sorted rows

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

1. Have a max heap with the first k elements of the matrix
2. Start traversing the remaining elements/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.

Edit:
You can sort the matrix rows based on their first element (break ties using the last element) in O(mlog(m)) time to reduce the number of heap push/pop operations. (For a monotonically decreasing matrix, you'll be pushing/popping a heap element at each case, which would degrade performance.)

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)+1, 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:

1. Scan the first column of matrix and find the minimum; call this low
2. Scan the last column of matrix and find the maximum; call this high
3. 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
4. Return low

The key to the solution is that all the rows are sorted. I will suggest the following:

Given a row-ordered matrix A:

1. Init a min heap with tuples of 0th index value of each row and their corresponding row numbers as min_heap where the location in the heap is defined by the value and the row number is for time optimization.
2. Init an index_vec = zeros([1, n]) to be the lngth equal to the number of rows where all the values in the vector are 0 since the starting index is 0 (n is the number of rows)
3. pop the top of min_heapand push into min_value_tuple
4. Advance index_vec[min_value_tuple[1]] += 1
5. push A[chosen_row, index_vec[min_value_tuple[1]]] into the min heap as a tuple with the corresponding row number.
6. return to step 3. for K times
7. pop the top of min_heap and return it

Note if during the loop in steps 3. through 6. index_vec[ii] = m where m is the number of colums, we will not push a new number to min_heap and just orgenize the heap to accomodate a new root.

Time complexity

1. Building the initial heaps takes O(n) since there are n numbers in the heap
2. K iteration of pushing numbers into the heap is O(klog(n))

Total runtime will be O(n + klong(n)).

Worst case is for k=m*n/2. In that case the time complexity will be O(n + nmlog(n)) = O(nmlog(n)).

Space

This requires a heap with size n and index_vec with size n as well so total space is O(n)

Note

If k > n*m/2 you can solve an equivalent problem if finding the n*m-k largest number in the matrix after flipping the rows of A.

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

Time complexity: (n + k) log n
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;
    }
}

Lets define a min-heap that stores a pair of values, then the algorithm consists of the following steps:

1- iterate through the values of the first column of the matrix and add each value to the heap as a pair of (value, row index).

2- repeat the following step exactly k-1 times: store the row index of the top pair of the heap in idx, then delete this pair. Now go to the row with index equal to idx and add the next unused element in it to the heap in the same way as before. After k-1 iterations the top element would contain the K-th element.

The time complexity of the first step is O(n.log n) such that n is the number of rows of the matrix. The time comexity of the second step is O(k.log n) because we will delet/add k-1 elements from/to the heap and at each moment the heap contains no more than n elements. So the total time complexity is O((n+k)log n).

Run quickselect on the first k elements of each row for O(kn) time.