LeetCode 633: Sum of Square Numbers

Problem Restatement

We are given a non-negative integer c.

We need to determine whether there exist two integers a and b such that:

a * a + b * b == c

If such integers exist, return True.

Otherwise, return False.

Input and Output

Example function shape:

def judgeSquareSum(c: int) -> bool:
    ...

Examples

Example 1:

c = 5

We can choose:

a = 1
b = 2

because:

1² + 2² = 1 + 4 = 5

So the answer is:

True

Example 2:

c = 3

Possible squares are:

No pair sums to 3.

So the answer is:

False

First Thought: Try Every Pair

A direct solution is:

1. Try every possible a.
2. Try every possible b.
3. Check whether:

a * a + b * b == c

The largest useful value is:

sqrt(c)

because larger squares already exceed c.

Code:

class Solution:
    def judgeSquareSum(self, c: int) -> bool:
        limit = int(c ** 0.5)

        for a in range(limit + 1):
            for b in range(limit + 1):
                if a * a + b * b == c:
                    return True

        return False

This works, but it takes:

O(c)

time in the worst case.

We can do better.

Key Insight

Suppose we fix one number a.

Then we need:

b² = c - a²

So after choosing a, we only need to know whether the remaining value is a perfect square.

Instead of checking every pair, we can use two pointers.

The smallest possible value is:

0

The largest useful value is:

floor(sqrt(c))

We maintain two pointers:

Then compute:

left² + right²

If the sum is too small, increase left.

If the sum is too large, decrease right.

This works because squares increase monotonically.

Algorithm

1. Set:

left = 0
right = floor(sqrt(c))

2. While left <= right:
   1. Compute:

current = left² + right²

2. If current == c, return True.

3. If current < c, increase left.

4. If current > c, decrease right.

5. If the loop finishes, return False.

Implementation

class Solution:
    def judgeSquareSum(self, c: int) -> bool:
        left = 0
        right = int(c ** 0.5)

        while left <= right:
            current = left * left + right * right

            if current == c:
                return True

            if current < c:
                left += 1
            else:
                right -= 1

        return False

Code Explanation

We start with the smallest and largest possible square roots:

left = 0
right = int(c ** 0.5)

At each step:

current = left * left + right * right

If the sum equals c, we found valid integers.

if current == c:
    return True

If the sum is too small:

if current < c:
    left += 1

we increase the smaller square.

If the sum is too large:

else:
    right -= 1

we decrease the larger square.

Eventually the pointers cross:

left > right

At that point, every possible pair has been checked.

Correctness

The algorithm maintains two integers left and right such that:

0 <= left <= right <= floor(sqrt(c))

Every possible valid pair (a, b) must lie within this range, because if either value were larger than sqrt(c), its square alone would exceed c.

At each step, the algorithm computes:

left² + right²

If this equals c, then the algorithm correctly returns True.

If the sum is smaller than c, then decreasing right would only make the sum smaller, because squares are non-decreasing. Therefore, the only possible way to reach c is to increase left.

If the sum is larger than c, then increasing left would only make the sum larger. Therefore, the only possible way to reach c is to decrease right.

Thus no valid pair is skipped.

The loop ends only after all feasible pairs have been eliminated. Therefore, if the algorithm returns False, no valid integers exist.

Complexity

Alternative Solution: Binary Search

Another approach is:

1. Iterate over all possible a.
2. Compute:

target = c - a²

3. Use binary search to check whether target is a perfect square.

class Solution:
    def judgeSquareSum(self, c: int) -> bool:
        limit = int(c ** 0.5)

        for a in range(limit + 1):
            target = c - a * a

            left = 0
            right = limit

            while left <= right:
                mid = (left + right) // 2
                square = mid * mid

                if square == target:
                    return True

                if square < target:
                    left = mid + 1
                else:
                    right = mid - 1

        return False

Complexity:

The two-pointer solution is simpler and faster.

Number Theory Insight

There is also a famous theorem:

A non-negative integer c can be written as the sum of two squares if and only if every prime factor of the form:

4k + 3

appears with an even exponent in the prime factorization of c.

Example:

5 = 1² + 2²

Prime factorization:

5

Since 5 is:

4k + 1

it is allowed.

But:

3

has factorization:

3

and 3 is:

4k + 3

with odd exponent 1, so it cannot be expressed as a sum of two squares.

This theorem gives another elegant mathematical solution.

Testing

def run_tests():
    s = Solution()

    assert s.judgeSquareSum(0) is True
    assert s.judgeSquareSum(1) is True
    assert s.judgeSquareSum(2) is True
    assert s.judgeSquareSum(3) is False
    assert s.judgeSquareSum(4) is True
    assert s.judgeSquareSum(5) is True
    assert s.judgeSquareSum(50) is True
    assert s.judgeSquareSum(99999999) is False

    print("all tests passed")

run_tests()

Test meaning: