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?

CodePudding user response：

If k is very large, We can make binary search on the answer, which leads to a solution with time complexity O(n*logN) where N is the range of each element, and n is the number of arrays.

What we need to learn is how to check some integer x whether <= correct answer or not. We can just enumerate each array, and make binary search on it to count the number of elements less than or equal to x. accumulate them, and compare it with k.

from typing import List
import bisect

def query_k_min(vecs: List[List[int]], k: int) -> int:
    # we assume each number >=1 and <=10^9
    l, r = 0, 10**9
    while r - l > 1:
        m = (l r)>>1
        tot = 0
        for vec in vecs:
            tot  = bisect.bisect_right(vec, m)
        if tot >= k: r = m
        else: l = m
    return r

a = [[2,3,6,7,9],[1,4,8,10],[2,3,5,7]]
for x in range(1,14):
    print(query_k_min(a,x))

CodePudding user response：

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.