Kth Smallest Element in multiple sorted arrays

Let’s say we have two arrays:

array1 = [2,3,6,7,9]

array2 = [1,4,8,10]

I understood how to find the kth element of two sorted arrays in log(min(m,n)) where m is the length of array1 and n is the length of array2 as follows:

def kthelement(arr1, arr2, m, n, k):
    if m > n:
        kthelement(arr2, arr1, n, m, k) 

    low = max(0, k - m)
    high = min(k, n)

    while low <= high:
        cut1 = (low + high) >> 1 
        cut2 = k - cut1 
        l1 = MIN_VALUE if cut1 == 0 else arr1[cut1 - 1] 
        l2 = MIN_VALUE if cut2 == 0 else arr2[cut2 - 1]
        r1 = MAX_VALUE if cut1 == n else arr1[cut1]
        r2 = MAX_VALUE if cut2 == m else arr2[cut2] 
        
        if l1 <= r2 and l2 <= r1:
            print(cut1, cut2)
            return max(l1, l2)
        elif l1 > r2:
            high = cut1 - 1
        else:
            low = cut1 + 1

But I couldn’t figure out how to extend this to multiple sorted arrays case. For example, given 3 arrays, I want to find the kth element of the final sorted array.

array1 = [2,3,6,7,9]

array2 = [1,4,8,10]

array3 = [2,3,5,7]

Is it possible to achieve it in log(min(m,n)) as in the two array case?

Answer

The general solution is to use a min-heap. If you have n sorted arrays and you want the kth smallest number, then the solution is O(k log n).

The idea is that you insert the first number from each array into the min-heap. When inserting into the heap, you insert a tuple that contains the number, and the array that it came from.

You then remove the smallest value from the heap and add the next number from the array that value came from. You do this k times to get the kth smallest number.

See https://www.geeksforgeeks.org/find-m-th-smallest-value-in-k-sorted-arrays/ for the general idea.