Leetcode: Square root of x

The task is to implement a square root method without using math.sqrt or ** or pow or any similar function. Any digits after the decimal point do not to be returned. This is effectively the floor(sqrt(n)) function. The first idea is to simply try out every number and multiply it with itself until the square root is found. If the next number is greater than n, return the previous number. This is however not very efficient as it has a time complexity of O(sqrt(n)) in every case. To speed this process up, we can use binary search. The idea is to always double the size of the steps until we reach the target number or overshoot. If we overshoot, take half of the number and start at 1 again, then double, and so on. The solution then looks as follows.

class Solution:
  def mySqrt(self, x: int) -> int:
    prefix = 0
    while True:
      i = 1
      while ((prefix + i) * (prefix + i)) <= x:
        i <<= 1
      prefix += i >> 1
      if prefix*prefix <= x and (prefix+1)*(prefix+1) > x:
        return prefix

if __name__ == '__main__':
  s = Solution().mySqrt
  print(s(125), 11)
  print(s(0), 0)
  print(s(4), 2)
  print(s(16), 4)