Find the median of two sorted arrays

median_of_two_sorted_arrays.py

# https://leetcode.com/problems/median-of-two-sorted-arrays/
from bisect import bisect_left
from itertools import takewhile


class Solution(object):

    def count_prefix(self, start_idx, st, char):
        """ Starting from the start index, checks how many times given character
        repeated"""
        return len(list(takewhile(lambda x: x == char, st[start_idx:])))

    def findMedianSortedArrays(self, nums1, nums2, median=None):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :rtype: float

        assume that nums1 has the value and search for value on num1
        if it is not in nums1 switch nums1 and num2 and restart search.
        """
        # make sure we have at least 1 element on num1
        if not nums1:
            if not nums2:
                return None
            else:
                return self.findMedianSortedArrays(nums2, nums1)

        # calculate median index
        # for total of 4 element median will be 1 (second element)
        # for total of 5 element median will be 2 (third element)
        median = (len(nums1)+len(nums2))
        odd = median % 2 == 1
        if odd:
            median = median + 1
        median = median/2 - 1
        start = -1
        end = len(nums1)
        while start+1 < end:
            middle = (end + start)/2
            # calculate number of extra values
            num2_index = bisect_left(nums2, nums1[middle])
            # calculate index of middle value in merged array
            merged_index = middle + num2_index
            repeated_chars = self.count_prefix(num2_index,
                                               nums2,
                                               nums1[middle])
            # if there are repeated number of elements,
            # we cannot exactly pinpoint the new index of the
            # element we choose. So we are checking a range
            # for given element.
            if abs(merged_index - median)<=repeated_chars:
                resp = nums1[middle]
                if not odd:
                    # if there is repeated characters
                    # we want to recalculate the num2_index to
                    # figureout what the next value will be.
                    num2_index = median - middle
                    next_item = min(nums1[middle+1:middle+2] +
                                    nums2[num2_index:num2_index+1])
                    resp = (resp + next_item)/2.0
                return resp
            elif merged_index > median:
                end = middle
            elif merged_index < median:
                start = middle
        # we could not find the index on nums1 switch nums1 and nums2 and
        # restart search.
        return self.findMedianSortedArrays(nums2, nums1)