initial tests completed

Author: koss-null

py3/median/median.py

@@ -1,69 +1,82 @@
 class Solution:
-    def single_place(self, nums, num):
-        """returns on which place >= num starts and how many same items"""
-        lf, rg = 0, len(nums) - 1
+    def bin_search(self, arr, num, fits=lambda a, b: a == b, le=lambda a, b: a <= b):
+        """
+        returns tuple: the leftmost place of num in arr and the amount of items same an num
+        """
+        lf, rg = 0, len(arr) - 1
         md = (lf + rg) // 2
-        while lf <= rg and nums[md] != num:
-            if nums[md] <= num:
+        while lf <= rg and not fits(arr[md], num):
+            if le(arr[md], num):
                 lf = md + 1
             else:
                 rg = md - 1
             md = (lf + rg) // 2
 
-        if nums[md] == num:
-            left_step, right_step = 1 >> 11, 1 >> 11
-            lf_sum, rg_sum = 0, 0
+        if fits(arr[md], num):
+            # searching the biggest possible step to each side and do it as many times as possible
+            left_step, right_step = 1 << 20, 1 << 20
+            lf_offset, rg_offset = 0, 0
             while left_step and right_step:
-                if md - left_step - lf_sum < 0 or nums[md - left_step - lf_sum] != num:
-                    left_step <<= 1
+                if md - left_step - lf_offset < 0 or not fits(arr[md - left_step - lf_offset], num):
+                    left_step >>= 1
                 else:
-                    lf_sum += left_step
+                    lf_offset += left_step
 
-                if md + right_step + rg_sum  >= len(nums) or nums[md + right_step + rg_sum] != num:
-                    right_step <<= 1
+                if md + right_step + rg_offset >= len(arr) or not fits(arr[md + right_step + rg_offset], num):
+                    right_step >>= 1
                 else:
-                    rg_sum += right_step
-
-            return [md - lf_sum, rg_sum + lf_sum + 1]
-        return [lf, 0]
+                    rg_offset += right_step
+            return [md - lf_offset, rg_offset + lf_offset, True]
+        return [lf, 0, False]
 
     def place(self, nums1, nums2, num):
         """returns the leftmost and the rightmost place of the element from nums1 in merged list"""
-        place1, place2 = self.single_place(nums1, num), self.single_place(nums2, num)
-        return [place1[0] + place2[0] + 2, place1[1] + place2[1]]
+        place1, place2 = self.bin_search(nums1, num), self.bin_search(nums2, num)
+        return [place1[0] + place2[0], place1[1] + int(place1[2]) + place2[1] + int(place2[2])]
 
     def find_place(self, nums, other_nums, desired_place):
-        lf, rg = 0, len(nums) - 1
-        md = (lf + rg) // 2
-        p = self.place(nums, other_nums, nums[md])
-        while lf <= rg and not (p[0] <= desired_place <= p[0] + p[1]):
-            if p[0] > desired_place:
-                rg = md - 1
-            else:
-                lf = md + 1
-            md = (lf + rg) // 2
-            p = self.place(nums, other_nums, nums[md])
-        return p
+        class Shared:
+            def __init__(self):
+                self.p = None
+
+            def fits(self, a, b):
+                self.p = Solution().place(nums, other_nums, a)
+                return self.p[0] <= b < self.p[0] + self.p[1]
+
+            def le(self, _, b):
+                return self.p[0] < b
+
+        s = Shared()
+        p = self.bin_search(
+            nums,
+            desired_place,
+            fits=s.fits,
+            le=s.le
+        )
+
+        if p[2]:
+            return p
+        return False
 
     def findMedianSortedArrays(self, nums1, nums2, desired_place=None, need_two=None):
         if desired_place is None:
-            desired_place = int((len(nums1) + len(nums2)) / 2.0 + 0.5)
+            desired_place = int((len(nums1) + len(nums2)) / 2.0 + 0.5) - 1
         if need_two is None:
             need_two = (len(nums1) + len(nums2)) % 2 == 0
 
         p = self.find_place(nums1, nums2, desired_place)
-        if p[0] <= desired_place < p[0] + p[1]:
-            second = nums1[p[0] - 1]
+        if p:
+            second = nums1[p[0]]
             if need_two:
                 second = self.findMedianSortedArrays(nums1, nums2, desired_place=desired_place+1, need_two=False)
-            return (nums1[p[0] - 1] + second) / 2.
+            return (nums1[p[0]] + second) / 2.
 
         p = self.find_place(nums2, nums1, desired_place)
-        if p[0] <= desired_place < p[0] + p[1]:
-            second = nums2[p[0] - 1]
+        if p:
+            second = nums2[p[0]]
             if need_two:
                 second = self.findMedianSortedArrays(nums1, nums2, desired_place=desired_place+1, need_two=False)
-            return (nums2[p[0] - 1] + second) / 2.
+            return (nums2[p[0]] + second) / 2.
 
         print("Something went wrong")
         return -1

py3/median/test/test1

@@ -1,4 +1,4 @@
 {
 "input": "1 3\n2",
-"output": "2.00000"
+"output": "2.0"
 }

py3/test.py

@@ -34,6 +34,7 @@ def check_result(test):
     out = test.output.decode("ascii").strip()
     if test.expected == out: 
         print_green(f"Test: {test.name} is [OK]")
+        print_green(out)
     else:
         print_red(f"Test: {test.name} is [FAIL]")
         print_green(test.expected)