square_root.py

class SquareRootException(Exception):
    pass

def square_root(n):
    if n < 0:
        raise SquareRootException(
            f'{n} is negative, so it does not have a square root.')
    if n == 0 or n == 1:
        return n
    start = 0
    end = n
    while True:
        guess = start + (end - start) // 2
        if end - start == 1:
            return guess
        if guess > n // guess:
            end = guess
        elif guess < n // guess:
            start = guess
        else:
            return guess

def square_root_recursive(n):
    if n < 0:
        raise SquareRootException(
            f'{n} is negative, so it does not have a square root.')
    if n == 0 or n == 1:
        return n
    return square_root_helper(n, 0, n)

def square_root_helper(n, start, end, depth=0):
    guess = start + (end - start) // 2
    if end - start == 1:
        print(f'Depth: {depth}')
        return guess
    if guess > n // guess:
        return square_root_helper(n, start, guess, depth + 1)
    elif guess < n // guess:
        return square_root_helper(n, guess, end, depth + 1)
    else:
        print(f'Depth: {depth}')
        return guess

def test_case(n, answer=None, exception=False):
    if not exception:
        guess = square_root(n)
        if answer == guess:
            print(f'CORRECT: {guess} is the floor(sqrt({n}))')
        else:
            print(f'ACTUAL: {guess} EXPECTED: {answer} is the floor(sqrt({n}))')
    else:
        try:
            guess = square_root(n)
            print(f'EXPECTED SquareRootException to be thrown, but got {guess}')
        except SquareRootException:
            print('CORRECT: Exception was thrown')

if __name__ == '__main__':
    # Simple integers with integer square roots
    test_case(4, answer=2)
    test_case(9, answer=3)
    test_case(16, answer=4)

    # These numbers don't have integer square roots, so should return
    # floor(sqrt(n))
    test_case(7, answer=2)
    test_case(11, answer=3)
    test_case(19, answer=4)

    # Base cases, nothing fancy is done for these values
    test_case(1, answer=1)
    test_case(0, answer=0)

    # Negative numbers should throw exception
    test_case(-1, exception=True)
    test_case(-4, exception=True)
    test_case(-11, exception=True)

    # This will trip recursion depth if you use recursion algorithm
    test_case(2 ** 8000, 2 ** 4000)