Find kth smallest element from two sorted arrays

Question

Given two sorted array of size M and N. I was trying to implement an algorithm with time complexity O(logM+logN). The strategy is basically comparing the middle index elements from two subarray based on their length condition.

    // Test case 1
    // Works for all position except when kth is 6
    int[] num1 = {6,7,8,9,10,11,12};
    int[] num2 = {1,2,3,4,5};

    // Test case 2
    // Always print the next smallest element
    int[] num3 = {1,3,5,7,9};
    int[] num4 = {2,4,6,8,10,12,14,16,18,20,22,24,30,40,50,56,77,35};


public static int findKth(int[] A, int p1, int r1, int[] B, int p2, int r2, int k){


    if (p1 > r1) { 
        return B[p2+k-1];
    } else if (p2 > r2) {
        return A[p1+k-1];
    }


    int midA = p1 + (int)Math.floor((r1-p1)/2);// Middle element from subarray A
    int midB = p2 + (int)Math.floor((r2-p2)/2);// Middle element from subarray B

    /**
     * Compare the sum of number of elements from left-subarray up to middle element. 
     */
    if ((midA-p1+midB-p2+2) < k) { 
        // We don't need to the left-subarray based on the comparisons between middle element
        if (A[midA] > B[midB]) {
            return findKth(A, p1, r1, B, midB+1, r2, k-(midB-p2+1)); //
        } else {
            return findKth(A, midA+1, r1, B, p2, r2, k-(midA-p1+1)); //
        }
    } else {
        // We don't need to the right-subarray based on the comparisons between middle element.
        if (A[midA] > B[midB]) {
            return findKth(A, p1, midA-1, B, p2, r2, k);
        } else {
            return findKth(A, p1, r1, B, p2, midB-1, k);
        }
    }
}

I felt the strategy I used should be correct. But for the two test cases shown above, it'll print wrong outputs in some particular kth value. So I guessed there must be something wrong with my strategy. Can anyone describe briefly which part of this implementation is not correct? Thanks!

Answer 1

I will update my answer if I can find bug in your snippet. Right now you can take a look of my code which logic is exactly same as yours except:

The strategy is basically comparing the middle index elements from two subarray based on their length condition.

The major difference for simplicity and small size of my code is I avoided some if-else conditions(for length condition) by calling the function with swapping its parameter in case of first array/index pair is not smaller.

public static int findKth(int[] A, int i, int[] B, int j, int k) {

    // Here is the simple trick. We've just changed the parameter order if first array is not smaller.
    // so that later we won't need to write if-else conditions to check smaller/greater stuff
    if((A.length - i) > (B.length - j))
    {
        return findKth(B, j, A, i, k);
    }

    if(i >= A.length) 
    {
        return B[j + k - 1];
    }
    if(k == 1)
    {
        return Math.min(A[i], B[j]);
    }

    int aMid = Math.min(k / 2, A.length - i);
    int bMid = k - aMid;

    if(A[i + aMid - 1] <= B[j + bMid - 1])
    {
        return findKth(A, i + aMid, B, j, k - aMid);
    }

    return findKth(A, i, B, j + bMid, k - bMid);
}

public static int findKthSmallestElement(int[] A, int[] B, int k)
{
    if(k > A.length + B.length)
        return -1;

    return findKth(A, 0, B, 0, k);
}

The time complexity is O(log(m + n)).

Answer 2

There are already good solutions to this problem on StackOverflow, for instance here and here, so I will focus on identifying the bug only:

This line:

if ((midA-p1+midB-p2+2) < k) { 

should be modified to use <=:

if ((midA-p1+midB-p2+2) <= k) { 

You can see why this is needed with this small example:

A = {1,3}
B = {2}

The other variables will be:

p1 = p2 = r2 = midA = midB = 0
r1 = 1

And so the expression midA-p1+midB-p2+2 will be 2. In this case, you clearly want to throw away an element of A, not the only element of B.

Note that in general it is OK to include the k case for the execution of the if block, as you will never throw away k (i.e. too many) elements there: the expression passed as last argument to the recursive call is always at the most k-1.

On the other hand, going to the else part when the if expression is k, is wrong, as midA or midB might well be the index of that k^th element, and it will be thrown out there.

Answer 3

#include <bits/stdc++.h>
using namespace std;

int findKthElement(int a[],int start1,int end1,int b[],int start2,int end2,int k){
    if(start1 >= end1)return b[start2+k-1];
    if(start2 >= end2)return a[start1+k-1];
    if(k==1)return min(a[start1],b[start2]);
    int aMax = INT_MAX;
    int bMax = INT_MAX;
    if(start1+k/2-1 < end1) aMax = a[start1 + k/2 - 1];
    if(start2+k/2-1 < end2) bMax = b[start2 + k/2 - 1];

    if(aMax > bMax){
        return findKthElement(a,start1,end1,b,start2+k/2,end2,k-k/2);
    }
    else{
        return findKthElement(a,start1 + k/2,end1,b,start2,end2,k-k/2);
    }
}

int main(void){
    int t;
    scanf("%d",&t);
    while(t--){
        int n,m,k;
        cout<<"Enter the size of 1st Array"<<endl;
        cin>>n;
        int arr[n];
        cout<<"Enter the Element of 1st Array"<<endl;
        for(int i = 0;i<n;i++){
            cin>>arr[i];
        }
        cout<<"Enter the size of 2nd Array"<<endl;
        cin>>m;
        int arr1[m];
        cout<<"Enter the Element of 2nd Array"<<endl;
        for(int i = 0;i<m;i++){
            cin>>arr1[i];
        }
        cout<<"Enter The Value of K";
        cin>>k;
        sort(arr,arr+n);
        sort(arr1,arr1+m);
        cout<<findKthElement(arr,0,n,arr1,0,m,k)<<endl;
    }

    return 0;
}

Time Complexcity is O(log(min(m,n)))