Square Functions

Every once in a while you come across a problem where you need a function that isn't in a standard Python Egg (library/package). What to do? Well, make your own of course!

I came across this problem on CodeFights:

Given an integer $n$, return the smallest integer $a$, such that $a$ is a square of some integer and is greater than $n$.

Example:

For $n = 5$, the output should be

nextSquare(n) = 9.

def intSqrt(n):
    x = n
    y = (x + 1) // 2
    while y < x:
        x = y
        y = (x + n // x) // 2
    return x

via Newton's method for solving $x^2-n=0$.

Let's check it's working:

intSqrt(10)

3

intSqrt(16)

4

Looks good!

Now let's get an 'isSquare' function that returns a Boolean.

def isSquare(n):
    if intSqrt(n)**2 == n:
        return True
    return False

Examples:

isSquare(10)

False

isSquare(36)

True

And finally:

def nextSquare(n):
    return (intSqrt(n)+1)**2

Does it work?

nextSquare(10)

16

nextSquare(16)

25

Sure, does! :)