Finding kth smallest number from n sorted arrays

Question

So, you have n sorted arrays (not necessarily of equal length), and you are to return the kth smallest element in the combined array (i.e the combined array formed by mergin

Answer 1

Old question, but none of the answers were good enough. So I am posting the solution using sliding window technique and heap:

class Node {

    int elementIndex;
    int arrayIndex;

    public Node(int elementIndex, int arrayIndex) {
        super();
        this.elementIndex = elementIndex;
        this.arrayIndex = arrayIndex;
    }

}

public class KthSmallestInMSortedArrays {

    public int findKthSmallest(List<Integer[]> lists, int k) {

        int ans = 0;
        PriorityQueue<Node> pq = new PriorityQueue<>((a, b) -> {
            return lists.get(a.arrayIndex)[a.elementIndex] -
                   lists.get(b.arrayIndex)[b.elementIndex];
        });

        for (int i = 0; i < lists.size(); i++) {
            Integer[] arr = lists.get(i);
            if (arr != null) {
                Node n = new Node(0, i);
                pq.add(n);
            }
        }

        int count = 0;

        while (!pq.isEmpty()) {
            Node curr = pq.poll();
            ans = lists.get(curr.arrayIndex)[curr.elementIndex];
            if (++count == k) {
                break;
            }

            curr.elementIndex++;
            pq.offer(curr);
        }

        return ans;
    }
}

The maximum number of elements that we need to access here is O(K) and there are M arrays. So the effective time complexity will be O(K*log(M)).

Answer 2

This could be considered the second half of a merge sort. We could simply merge all the sorted lists into a single list...but only keep k elements in the combined lists from merge to merge. This has the advantage of only using O(k) space, but something slightly better than merge sort's O(n log n) complexity. That is, it should in practice operate slightly faster than a merge sort. Choosing the kth smallest from the final combined list is O(1). This is kind of complexity is not so bad.