LeetCode 4. Median of Two Sorted Arrays

median-of-two-sorted-arrays.java

/**
 * https://leetcode.com/problems/median-of-two-sorted-arrays/
 *
 * O(log (m + n)) runtime, O(1) space
 * where m is the length of the first array, and n is the length of the second array
 *
 * this question can be solved by using the solution of Find the k-th Smallest Element in the Union of Two Sorted Arrays
 *
 * 1. if m + n is odd, the median would be the ((m + n) / 2 + 1)-th smallest element
 * 2. otherwise, the median would be the mean of
 *    the ((m + n) / 2)-th smallest element and the ((m + n) / 2 + 1)-th smallest element
 */

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int len = nums1.length + nums2.length;
        if (len % 2 == 1) {
            return findKth(nums1, 0, nums2, 0, len / 2 + 1);
        } else {
            return ((double) findKth(nums1, 0, nums2, 0, len / 2) +
                    (double) findKth(nums1, 0, nums2, 0, len / 2 + 1)
                   ) / 2; // note that it's better than
                          // (findKth(nums1, 0, nums2, 0, len / 2) + findKth(nums1, 0, nums2, 0, len / 2 + 1)) / 2
                          // because of the overflow error
        }
    }
    
    // find the k-th smallest element of two sorted arrays
    private int findKth(int[] nums1, int start1, int[] nums2, int start2, int k) {
        // note the first 2 if statements here
        if (start1 == nums1.length) { // ignore the first array
            return nums2[start2 + k - 1];
        }
        if (start2 == nums2.length) { // ignore the second array
            return nums1[start1 + k - 1];
        }
        // now take both arrays into consideration
        if (k == 1) {
            return Math.min(nums1[start1], nums2[start2]);
        }
        // skip k / 2 elements every time
        if (start1 + k / 2 - 1 >= nums1.length) { // less than k / 2 elements in the first array
            return findKth(nums1, start1, nums2, start2 + k / 2, k - k / 2);
        }
        if (start2 + k / 2 - 1 >= nums2.length) { // less than k / 2 elements in the second array
            return findKth(nums1, start1 + k / 2, nums2, start2, k - k / 2);
        }
        // now the number of elements of both arrays is greater than or equal to k / 2
        int key1 = nums1[start1 + k / 2 - 1];
        int key2 = nums2[start2 + k / 2 - 1];
        if (key1 < key2) {
            return findKth(nums1, start1 + k / 2, nums2, start2, k - k / 2);
        } else {
            return findKth(nums1, start1, nums2, start2 + k / 2, k - k / 2);
        }
    }
}