Sum of Square Numbers — Medium Problem & Solution

Given a non-negative integer c, decide whether there exist non-negative integers a and b with a² + b² == c. Return true if such a pair exists.

Difficulty: Medium
Topics: Math, Two Pointers, Binary Search
Asked at: Amazon, Google, Adobe
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

Given a non-negative integer c, decide whether there exist non-negative integers a and b with a² + b² == c.

Return true if such a pair exists.

Example 1

Input: c = 5
Output: true
Explanation: 1² + 2² = 5.

Example 2

Input: c = 3
Output: false

Example 3

Input: c = 4
Output: true
Explanation: 0² + 2² = 4.

Constraints

0 <= c <= 10000000

How to solve Sum of Square Numbers

Search the pair (a, b) with a <= b <= sqrt(c) using two pointers. The sum a² + b² rises when a rises and falls when b falls, so every candidate is covered in one sweep.

Approach

Find the largest b with b² <= c by an integer loop.
Start a at 0 and compare a² + b² with c.
Increase a when the sum is too small; decrease b when it is too large; report success on equality.
Stop when the pointers cross.

Why it works

Any solution has a <= b after swapping, and both are at most sqrt(c). The sweep is monotone in both directions, so it neither skips a solution nor revisits a pair.

Complexity

Time — O(sqrt(c))
Space — O(1)

Pitfalls

A floating-point sqrt can land one off on a perfect square, and the judge's C harness has no math.h — the integer loop avoids both.
a and b may be zero, so c = 0 and c = 4 both answer true.
b * b stays inside 32 bits at this limit, but at LeetCode's real bound it needs a wider type.

Reference solution

Python

def judgeSquareSum(c: int) -> bool:
    a, b = 0, 0
    while (b + 1) * (b + 1) <= c:
        b += 1
    while a <= b:
        total = a * a + b * b
        if total == c:
            return True
        if total < c:
            a += 1
        else:
            b -= 1
    return False

JavaScript

var judgeSquareSum = function(c) {
    var a = 0, b = 0;
    while ((b + 1) * (b + 1) <= c) b++;
    while (a <= b) {
        var sum = a * a + b * b;
        if (sum === c) return true;
        if (sum < c) a++;
        else b--;
    }
    return false;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 213 math problems · the whole catalogue