Median of Two Sorted Arrays 解答

Question

There are two sorted arrays nums1 and nums2 of size m and n respectively. Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).

Solution 1 -- Traverse Array

Use merge procedure of merge sort here. Keep track of count while comparing elements of two arrays. Note to consider odd / even situation.

Time complexity O(n), space cost O(1).

 public class Solution {
     public double findMedianSortedArrays(int[] nums1, int[] nums2) {
         int length1 = nums1.length, length2 = nums2.length, total = length1 + length2;
         int index1 = total / 2, index2;
         // Consider two situations: (n + m) is odd, (n + m) is even
         if (total % 2 == 0)
             index2 = total / 2 - 1;
         else
             index2 = total / 2;
         // Traverse once to get median
         int p1 = 0, p2 = 0, p = -1, tmp, median1 = 0, median2 = 0;
         while (p1 < length1 && p2 < length2) {
             if (nums1[p1] < nums2[p2]) {
                 tmp = nums1[p1];
                 p1++;
             } else {
                 tmp = nums2[p2];
                 p2++;
             }
             p++;
             if (p == index1)
                 median1 = tmp;
             if (p == index2)
                 median2 = tmp;
         }
         if (p < index1 || p < index2) {
             while (p1 < length1) {
                 tmp = nums1[p1];
                 p1++;
                 p++;
                 if (p == index1)
                     median1 = tmp;
                 if (p == index2)
                     median2 = tmp;
             }
             while (p2 < length2) {
                 tmp = nums2[p2];
                 p2++;
                 p++;
                 if (p == index1)
                     median1 = tmp;
                 if (p == index2)
                     median2 = tmp;
             }
         }
         return ((double)median1 + (double)median2) / 2;
     }
 }

Solution 2 -- Binary Search

General way to find Kth element in two sorted arrays. Time complexity O(log(n + m)).

 public class Solution {
     public double findMedianSortedArrays(int[] nums1, int[] nums2) {
         int m = nums1.length, n = nums2.length;
         if ((m + n) %2 == 1)
             return findKthElement(nums1, nums2, (m + n) / 2, 0, m - 1, 0, n - 1);
         else
             return (findKthElement(nums1, nums2, (m + n) / 2, 0, m - 1, 0, n - 1) * 0.5 +
                     findKthElement(nums1, nums2, (m + n) / 2 - 1, 0, m - 1, 0, n - 1) * 0.5);
     }

     private double findKthElement(int[] A, int[] B, int k, int startA, int endA, int startB, int endB) {
         int l1 = endA - startA + 1;
         int l2 = endB - startB + 1;
         if (l1 == 0)
             return B[k + startB];
         if (l2 == 0)
             return A[k + startA];
         if (k == 0)
             return A[startA] > B[startB] ? B[startB] : A[startA];
         int midA = k * l1 / (l1 + l2);
         // Note here
         int midB = k - midA - 1;
         midA = midA + startA;
         midB = midB + startB;
         if (A[midA] < B[midB]) {
             k = k - (midA - startA + 1);
             endB = midB;
             startA = midA + 1;
             return findKthElement(A, B, k, startA, endA, startB, endB);
         } else if (A[midA] > B[midB]) {
             k = k - (midB - startB + 1);
             endA = midA;
             startB = midB + 1;
             return findKthElement(A, B, k, startA, endA, startB, endB);
         } else {
             return A[midA];
         }
     }
 }