algorithms/findMedianSortedArrays.py at master · fafk/algorithms · GitHub

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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""
Given two sorted arrays, find their median.

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two
sorted arrays.

Follow up: The overall run time complexity should be O(log (m+n)).


Example 1:

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.

Example 2:

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.

Example 3:

Input: nums1 = [0,0], nums2 = [0,0]
Output: 0.00000

Example 4:

Input: nums1 = [], nums2 = [1]
Output: 1.00000

Example 5:

Input: nums1 = [2], nums2 = []
Output: 2.00000

"""
from typing import List
from math import floor


class Solution:
    # TODO the alg can stop after it's readched (m+n)/2 in the new array
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        assert len(nums1) and len(nums2)
        i = 0
        j = 0
        out = []
        while not (i == len(nums1) and j == len(nums2)):
            if i == len(nums1):
                out.append(nums2[j])
                j = j + 1
                continue
            if j == len(nums2):
                out.append(nums1[i])
                i = 1 + i
                continue

            if nums1[i] < nums2[j]:
                out.append(nums1[i])
                i = i + 1
            else:
                out.append(nums2[j])
                j = j + 1

        if len(out) % 2 == 1:
            return out[floor(len(out) / 2)]
        else:
            return (out[floor(len(out) / 2)] + out[floor(len(out) / 2) - 1]) / 2

    def easy_but_slow(self, nums1: List[int], nums2: List[int]) -> float:
        assert len(nums1) and len(nums2)
        out = sorted(nums1 + nums2)
        if len(out) % 2 == 1:
            return out[floor(len(out) / 2)]
        else:
            return (out[floor(len(out) / 2)] + out[floor(len(out) / 2) - 1]) / 2


assert Solution().findMedianSortedArrays([1, 3, 4], [2, 2, 10]) == 2.5 == Solution().easy_but_slow(
    [1, 3, 4], [2, 2, 10])