-
Notifications
You must be signed in to change notification settings - Fork 0
/
002 Median of Two Sorted Arrays.py
executable file
·65 lines (53 loc) · 2.15 KB
/
002 Median of Two Sorted Arrays.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
"""
There are two sorted arrays A and B 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)).
"""
__author__ = 'Danyang'
class Solution:
def findMedianSortedArrays(self, A, B):
"""
Merge two arrays to get the median, O((m+n)/2)
Algorithm: Find k-th element in 2 array
A: A_left A[m/2] A_right
B: B_left B[n/2] A_right
if A[m/2]>B[n/2] and k>m/2+n/2, then disregard B_left and B[n/2]
if A[m/2]>B[n/2] and k<=m/2+n/2, then disregard A_right and A[m/2]
if A[m/2]<=B[n/2] and k>m/2+n/2, then disregard A_left and A[m/2]
if A[m/2]<=B[n/2] and k<=m/2+n/2, then disregard B_right and B[n/2]
whether to disregard A[m/2] or B[n/2] takes time to consider
T(N) = T(3/4 N) + O(1), thus T(N) = O(lg N) where N = |A|+|B|
O(log (m+n)), thus binary search.
:param A: list
:param B: list
:return: float
"""
m = len(A)
n = len(B)
if ((m+n)&1 == 0):
return (self.find_kth(A, B, (m+n)/2)+self.find_kth(A, B, (m+n)/2-1))/2.0
else:
return self.find_kth(A, B, (m+n)/2)
def find_kth(self, A, B, k):
"""
:param A:
:param B:
:param k: index starting from 0
:return:
"""
if not A: return B[k]
if not B: return A[k]
if k == 0: return min(A[0], B[0])
m, n = len(A), len(B)
# pay attention to consider the equal sign. Assigning equal sign is an art.
if A[m/2] >= B[n/2]:
if k > m/2+n/2:
return self.find_kth(A, B[n/2+1:], k-n/2-1) # exclude B[n/2] to make progress
else:
return self.find_kth(A[:m/2], B, k) # exclude A[m/2] to make progress
else:
return self.find_kth(B, A, k)
if __name__ == "__main__":
assert Solution().findMedianSortedArrays([1, 2], [1, 2, 3]) == 2
assert Solution().findMedianSortedArrays([1, 2], [3]) == 2
assert Solution().findMedianSortedArrays([1], [2, 3]) == 2
assert Solution().findMedianSortedArrays([1, 2], [1, 2]) == 1.5