Coverage for blog/dsa/leetcode/median_two_sorted_arrays/__init__.py: 62%

1import pytest 

2 

3 

4# start snippet solution_initial 

5class Solution: 

6 

7 # NOTE: Input arrays should already be sorted. 

8 def findMedianSortedArrays( 

9 self, 

10 nums1: list[int], 

11 nums2: list[int], 

12 ) -> float: 

13 

14 n1, n2 = len(nums1), len(nums2) 

15 if n1 + n2 == 1: 

16 if n1 != 0: 

17 return nums1[0] 

18 else: 

19 return nums2[0] 

20 

21 start_1, start_2 = 0, 0 

22 stop_1, stop_2 = n1 - 1, n2 - 1 

23 left, right = None, None 

24 

25 while True: 

26 # If both, decide which one to decrement/increment 

27 if start_1 <= stop_1 and start_2 <= stop_2: 

28 a, b = nums1[start_1], nums2[start_2] 

29 if a <= b: 

30 left = a 

31 start_1 += 1 

32 else: 

33 left = b 

34 start_2 += 1 

35 

36 a, b = nums1[stop_1], nums2[stop_2] 

37 if a >= b: 

38 right = a 

39 stop_1 -= 1 

40 else: 

41 right = b 

42 stop_2 -= 1 

43 elif start_1 < stop_1: 

44 left = nums1[start_1] 

45 right = nums1[stop_1] 

46 

47 stop_1 -= 1 

48 start_1 += 1 

49 elif start_2 < stop_2: 

50 left = nums2[start_2] 

51 right = nums2[stop_2] 

52 

53 stop_2 -= 1 

54 start_2 += 1 

55 elif stop_1 == start_1: 

56 left = right = nums1[stop_1] 

57 break 

58 elif stop_2 == start_2: 

59 left = right = nums2[stop_2] 

60 break 

61 else: 

62 break 

63 

64 return (left + right) / 2 # type: ignore[operator] 

65 

66 

67# end snippet solution_initial 

68 

69 

70# start snippet solution_2 

71class Solution2: 

72 

73 def findMedianSortedArrays( 

74 self, 

75 nums1: list[int], 

76 nums2: list[int], 

77 ) -> float: 

78 

79 n1, n2 = len(nums1), len(nums2) 

80 start_1, start_2 = 0, 0 

81 stop_1, stop_2 = n1 - 1, n2 - 1 

82 left, right = None, None 

83 

84 while True: 

85 # If both, decide which one to decrement/increment 

86 if start_1 <= stop_1 and start_2 <= stop_2: 

87 a, b = nums1[start_1], nums2[start_2] 

88 if a <= b: 

89 left = a 

90 start_1 += 1 

91 else: 

92 left = b 

93 start_2 += 1 

94 

95 a, b = nums1[stop_1], nums2[stop_2] 

96 if a >= b: 

97 right = a 

98 stop_1 -= 1 

99 else: 

100 right = b 

101 stop_2 -= 1 

102 

103 # If one remains, find its median 

104 elif start_1 <= stop_1: 

105 left = nums1[start_1] 

106 right = nums1[stop_1] 

107 

108 stop_1 -= 1 

109 start_1 += 1 

110 elif start_2 <= stop_2: 

111 left = nums2[start_2] 

112 right = nums2[stop_2] 

113 

114 stop_2 -= 1 

115 start_2 += 1 

116 else: 

117 break 

118 

119 return (left + right) / 2 # type: ignore[operator] 

120 

121 

122# end snippet solution_2 

123 

124 

125@pytest.fixture 

126def solution(): 

127 return Solution() 

128 

129 

130# 1,2,3, 

131# 4,5,6, 

132# 7,8,9, 

133# 10,11, 12, 

134# 13,14,15, 

135# 16,17 

136 

137 

138@pytest.mark.parametrize( 

139 "nums1, nums2, answer", 

140 ( 

141 ([3], [-2, -1], -1), 

142 ([1], [], 1), 

143 ([1, 2, 3], [], 2), 

144 ([1, 3], [2], 2.0), 

145 ([1, 2], [3, 4], 2.5), 

146 ([5, 6, 7, 8], [4, 9], 6.5), 

147 ([1, 2, 5, 6], [3, 4], 3.5), 

148 ([1, 2, 3, 7, 8, 9], [4, 5, 6], 5.0), 

149 ([2, 2, 4, 4], [2, 2, 2, 4, 4], 2.0), 

150 ([1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17], 9.0), 

151 ), 

152) 

153def test_solution( 

154 solution: Solution, 

155 nums1: list[int], 

156 nums2: list[int], 

157 answer: float, 

158): 

159 assert solution.findMedianSortedArrays(nums1, nums2) == answer