Python Program for Median of two sorted arrays of same size

# A Simple Merge based O(n) Python 3 solution 
# to find median of two sorted lists

# This function returns median of ar1[] and ar2[].
# Assumptions in this function:
# Both ar1[] and ar2[] are sorted arrays
# Both have n elements
def getMedian( ar1, ar2 , n):
    i = 0 # Current index of i/p list ar1[]
    
    j = 0 # Current index of i/p list ar2[]
    
    m1 = -1
    m2 = -1
    
    # Since there are 2n elements, median
    # will be average of elements at index
    # n-1 and n in the array obtained after
    # merging ar1 and ar2
    count = 0
    while count < n + 1:
        count += 1
        
        # Below is to handle case where all
        # elements of ar1[] are smaller than
        # smallest(or first) element of ar2[]
        if i == n:
            m1 = m2
            m2 = ar2[0]
            break
        
        # Below is to handle case where all 
        # elements of ar2[] are smaller than
        # smallest(or first) element of ar1[]
        elif j == n:
            m1 = m2
            m2 = ar1[0]
            break
        # equals sign because if two 
        # arrays have some common elements 
        if ar1[i] <= ar2[j]:
            m1 = m2 # Store the prev median
            m2 = ar1[i]
            i += 1
        else:
            m1 = m2 # Store the prev median
            m2 = ar2[j]
            j += 1
    return (m1 + m2)/2

# Driver code to test above function
ar1 = [1, 12, 15, 26, 38]
ar2 = [2, 13, 17, 30, 45]
n1 = len(ar1)
n2 = len(ar2)
if n1 == n2:
    print("Median is ", getMedian(ar1, ar2, n1))
else:
    print("Doesn't work for arrays of unequal size")

# This code is contributed by "Sharad_Bhardwaj".

# Python program for above approach

# function to return median of the arrays
# both are sorted & of same size
def getMedian(ar1, ar2, n):
    i, j = n - 1, 0

    # while loop to swap all smaller numbers to arr1
    while(ar1[i] > ar2[j] and i > -1 and j < n):
        ar1[i], ar2[j] = ar2[j], ar1[i]
        i -= 1
        j += 1

    ar1.sort()
    ar2.sort()

    return (ar1[-1] + ar2[0]) >> 1


# Driver program
if __name__ == '__main__':
    ar1 = [1, 12, 15, 26, 38]
    ar2 = [2, 13, 17, 30, 45]

    n1, n2 = len(ar1), len(ar2)

    if(n1 == n2):
        print('Median is', getMedian(ar1, ar2, n1))
    else:
        print("Doesn't work for arrays of unequal size")

# This code is contributed by saitejagampala

# Calculate the number of elements less than or equal to mid in the given arrays
def count_less_than_or_equal_to_mid(mid, arrays):
    count = 0
    for array in arrays:
        count += len([x for x in array if x <= mid])
    return count

def find_kth_element(arrays, n):
    ans = 0.0
    low = -1e9
    high = 1e9
    pos = n

    # Binary search to find the kth element
    while low <= high:
        mid = low + (high - low) // 2
        count = count_less_than_or_equal_to_mid(mid, arrays)
        if count <= pos:
            low = mid + 1
        else:
            high = mid - 1

    ans = low

    # Update pos and repeat the binary search to find the (n-1)th element
    pos = n - 1
    low = -1e9
    high = 1e9
    while low <= high:
        mid = low + (high - low) // 2
        count = count_less_than_or_equal_to_mid(mid, arrays)
        if count <= pos:
            low = mid + 1
        else:
            high = mid - 1

    ans += low

    # Return the average of the two elements
    return (ans / 2.0)

# Test with some arrays
arrays = [[1, 4, 5, 6, 10], [2, 3, 4, 5, 7]]
n = 5
print("Median in", find_kth_element(arrays, n)) 

#code is contributed by khushboogoyal499