Search code examples
pythonalgorithmbinary-searchmedian

Finding median of two sorted arrays. Can some inequality checks be eliminated?


Working on this problem and post code, my question is whether it is safe to change this line of code

j > 0 and i < m and B[j-1] > A[i]

to

i < m and B[j-1] > A[i]

and also it is safe to change this line of code

i > 0 and j < n and A[i-1] > B[j]

to

i > 0 and A[i-1] > B[j]

I think remove the condition check of j is safe since we already making sure size of A is no bigger than size of B.

Problem statement

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)).

Implementation

 def median(A, B):
    m, n = len(A), len(B)
    if m > n:
        A, B, m, n = B, A, n, m
    if n == 0:
        raise ValueError

    imin, imax, half_len = 0, m, (m + n + 1) / 2
    while imin <= imax:
        i = (imin + imax) / 2
        j = half_len - i
        if j > 0 and i < m and B[j-1] > A[i]:
            # i is too small, must increase it
            imin = i + 1
        elif i > 0 and j < n and A[i-1] > B[j]:
            # i is too big, must decrease it
            imax = i - 1
        else:
            # i is perfect
            if i == 0: max_of_left = B[j-1]
            elif j == 0: max_of_left = A[i-1]
            else: max_of_left = max(A[i-1], B[j-1])

            if (m + n) % 2 == 1:
                return max_of_left

            if i == m: min_of_right = B[j]
            elif j == n: min_of_right = A[i]
            else: min_of_right = min(A[i], B[j])

            return (max_of_left + min_of_right) / 2.0

Solution

  • Yes I think, you can remove the condition j > 0, because

    j = half_len - i
    

    and you already check that i<m and (m + n + 1) / 2 must be bigger than m since n>=m

    same for the second condition j < n. You already make sure that i>0, which ensures that j can at most be (2n+1)/2 - 1 which is smaller than n and thus automatically satisfies your condition