from binary_search import binarySearchFirst, binarySearchLast

import math
import random

class Solution:
def findMedianSortedArrays(self, nums1, nums2):
n1 = len(nums1)
n2 = len(nums2)

self.nums1 = nums1
self.nums2 = nums2

if not nums1:
return self.find_median_one_array(nums2, 0, n2)
if not nums2:
return self.find_median_one_array(nums1, 0, n1)

return self.find_median_2(0, n1, 0, n2, (n1 + n2) // 2)

def find_median_2(self, start1, end1, start2, end2, residue):

c = random.choice([0, 1])
if start1 == end1:
c = 1
if start2 == end2:
c = 0
if start1 == start2 == end1 == end2:
raise Exception("awet")
if c == 0:
m = self.find_median_one_array(self.nums1, start1, end1)
if start2 == end2:
ins1, ins2 = start2, start2
else:
ins1, ins2 = self.find_insertion_range(self.nums2, m, start2, end2)

# z = [(end1 - start1) // 2 + (ins - start2) for ins in range(ins1,

ins2+1)]

z1 = (end1 - start1) // 2 + (ins1 - start2)

z2 = (end1 - start1) // 2 + (ins2 - start2)

print("H1 ", start1, end1, m, z1, z2, residue)

if z1 <= residue <= z2:

ins = ins1 + residue - z1

if (end1 - start1) % 2 == 0:
if start1 + (end1-start1) //2 - 1 < 0:
prev1 = - math.inf
else:
prev1 = self.nums1[start1 + (end1-start1) //2 - 1]
if start1 + (end1-start1)//2 >= len(self.nums1):
next1 = math.inf
else:
next1 = self.nums1[start1 + (end1-start1)//2]
else:
if math.floor(start1 + (end1-start1)/2) -1 < 0:
prev1 = - math.inf
else:
prev1 = self.nums1[start1 + math.floor((end1-start1)/2) -
1]
if math.ceil(start1 + (end1-start1)/2) >= len(self.nums1):
next1 = math.inf
else:
next1 = self.nums1[start1 + math.ceil((end1-start1)/2)]
if ins - 1 < 0:
prev2 = - math.inf
else:
prev2 = self.nums2[ins-1]
if ins >= len(self.nums2):
next2 = math.inf
else:
next2 = self.nums2[ins]

print("1 ", prev1, prev2, next1, next2, m)

p = max(prev1, prev2)
n = min(next1, next2)

if (len(self.nums1) + len(self.nums2)) % 2 == 0:
if (end1 - start1) % 2 == 0:
return (p + n) / 2
else:
return (m + p) / 2
else:
if (end1 - start1) % 2 == 0:
return n
else:
return m
elif z2 < residue:
return self.find_median_2(start1 + (end1 - start1) // 2, end1,
ins2, end2, residue - z2)
else:
return self.find_median_2(start1, start1 + (end1 - start1) // 2,
start2, ins1, residue)
else:
m = self.find_median_one_array(self.nums2, start2, end2)
if start1 == end1:
ins1, ins2 = start1, start1
else:
ins1, ins2 = self.find_insertion_range(self.nums1, m, start1, end1)

z1 = (end2 - start2) // 2 + (ins1 - start1)

z2 = (end2 - start2) // 2 + (ins2 - start1)

print("H2 ", start2, end2, m, z1, z2, residue)

if z1 <= residue <= z2:

ins = ins1 + residue - z1

if (end2 - start2) % 2 == 0:
if start2 + (end2-start2) //2 - 1 < 0:
prev1 = - math.inf
else:
prev1 = self.nums2[start2 + (end2-start2) // 2 - 1]
if start2 + (end2-start2)//2 >= len(self.nums2):
next1 = math.inf
else:
next1 = self.nums2[start2 + (end2-start2)// 2]
 else:
if math.floor(start2 + (end2-start2)/2) - 1< 0:
prev1 = - math.inf
else:
prev1 = self.nums2[start2 + math.floor((end2-start2)/2)-1]
if math.ceil(start2 + (end2-start2)/2) >= len(self.nums2):
next1 = math.inf
else:
next1 = self.nums2[start2 + math.ceil((end2-start2)/2)]
if ins - 1 < 0:
prev2 = - math.inf
else:
prev2 = self.nums1[ins-1]
if ins >= len(self.nums1):
next2 = math.inf
else:
next2 = self.nums1[ins]

print("2 ", prev1, prev2, next1, next2, m)

p = max(prev1, prev2)
n = min(next1, next2)

if (len(self.nums1) + len(self.nums2)) % 2 == 0:
if (end2 - start2) % 2 == 0:
return (p + n) / 2
else:
return (m + p) / 2
else:
if (end2 - start2) % 2 == 0:
return n
else:
return m
elif z2 < residue:
return self.find_median_2(ins2, end1, start2 + (end2-start2)//2,
end2, residue - z2)
else:
return self.find_median_2(start1, ins1, start2, start2 + (end2 -
start2) // 2, residue)

def find_median_two_arrays(self, start1, end1, start2, end2, residue):

m = self.find_median_one_array(self.nums1, start1, end1)
# ins = self.find_insertion_index(self.nums2, m, start2, end2)
ins1, ins2 = self.find_insertion_range(self.nums2, m, start2, end2)

z = [(end1 - start1) // 2 + (ins - start2) for ins in range(ins1, ins2+1)]

# print(m, ins, z)
# raise Exception("Stop")

if residue in z:
return m
elif z[-1] < residue:
return self.find_median_two_arrays((end1 - start1) // 2, end1, ins2,
end2, residue - z)
else:
return self.find_median_two_arrays(start1, (end1 - start1) // 2,
start2, ins1, residue)

def find_insertion_range(self, nums, key, start, end):

idx1 = binarySearchFirst(nums, key, start, end - 1)

try:
idx2 = binarySearchLast(nums, key, start, end - 1)
except Exception:
print(nums)
print(key, start, end)
import sys
sys.exit()

return idx1, idx2

def find_insertion_index(self, nums, key, start, end):

for i, num in enumerate(nums[start:end]):
if num <= key:
continue
else:
return i
# raise Exception("Insertion Error")
return i + 1

def find_median_one_array(self, nums, start, end):

n = len(nums[start:end])
i = start + n // 2
if n % 2 == 0:
try:
t = (nums[i-1] + nums[i]) / 2
except Exception:
print(nums, start, end, i)
import sys
sys.exit()
return (nums[i-1] + nums[i]) / 2
else:
return nums[i]

if __name__ == "__main__":
a = [1,2]
b = [3,4]
s = Solution()
# print(s.findMedianSortedArrays(a, b))
n = 10000
for i in range(n):
a = list(sorted([random.randint(-5,5) for j in range(22)]))
b = list(sorted([random.randint(-5,5) for j in range(32)]))
ab = list(sorted(a+b))
m1 = s.find_median_one_array(ab, 0, len(ab))
m2 = s.findMedianSortedArrays(a, b)
assert m1 == m2, a + b
# print(s.find_insertion_range([3,4],1.5, 0, 2))

a = [-2]
b = [-3, -2]
print(s.findMedianSortedArrays(a,b))
ab = list(sorted(a+b))
print(s.find_median_one_array(ab, 0, len(ab)))