Kth element of 2 sorted arrays

Kth element of 2 sorted arrays

Photo by Jonny Clow on Unsplash

Problem

Given two sorted arrays arr1 and arr2 of size N and M respectively and an element K. The task is to find the element that would be at the kth position of the final sorted array.

Example 1:

Input:
arr1[] = {2, 3, 6, 7, 9}
arr2[] = {1, 4, 8, 10}
k = 5
Output:
6
Explanation:
The final sorted array would be -
1, 2, 3, 4, 6, 7, 8, 9, 10
The 5th element of this array is 6.

Example 2:

Input:
arr1[] = {100, 112, 256, 349, 770}
arr2[] = {72, 86, 113, 119, 265, 445, 892}
k = 7
Output:
256
Explanation:
Final sorted array is - 72, 86, 100, 112,
113, 119, 256, 265, 349, 445, 770, 892
7th element of this array is 256.

Constraints:
1 <= N, M <= 106
0 <= arr1i, arr2i < INT_MAX
1 <= K <= N+M

Solution

This approach is exactly same as the median of two sorted arrays (link). The only thing that changes is the number of elements on the left, which will be k. If we find that the array is sorted, then we return the maximum of l1 and l2.

Also,

  • low = max(0, k - m): The reason is if there are fewer than k elements in the second array, then it might happen that we need at least k-m elements from the first array. Also, if we take low as 0, then it might happen that mid1 also takes 0. In that case, mid2 will be k, and if the length of the second array is less than k, then it will give an out-of-bounds exception.

  • high = min(n, k): If k < n, then it cannot point to n because mid2 will give a negative index.

To understand the starting values of low and high in-depth, directly substitute the extreme values in mid, then check whether it is negative or out of bounds for the given array.

class Solution {
    public long kthElement( int arr1[], int arr2[], int n, int m, int k) {
        if(n>m) return kthElement(arr2, arr1, m, n, k);
        int low = Math.max(0,k-m);
        int high = Math.min(n,k);

        while(low<=high){
            int mid1 = (low+high)/2;
            int mid2 = k - mid1;

            int l1 = Integer.MIN_VALUE;
            int l2 = Integer.MIN_VALUE;
            int r1 = Integer.MAX_VALUE;
            int r2 = Integer.MAX_VALUE;

            if(mid1-1>=0) l1 = arr1[mid1-1];
            if(mid2-1>=0) l2 = arr2[mid2-1];
            if(mid1<n) r1 = arr1[mid1];
            if(mid2<m) r2 = arr2[mid2];

            if(l1<=r2 && l2<=r1) return Math.max(l1,l2);

            if(l1>r2) {
                high = mid1 -1;
            }
            else if(l2>r1){
                low = mid1 +1;
            }
        }
        return -1; 
    }
}