633. Sum of Square Numbers

Problem

Given a non-negative integer c, decide whether there’re two integers a and b such that a² + b² = c.

https://leetcode.cn/problems/sum-of-square-numbers/

Example 1:

Input: c = 5
Output: true
Explanation: 1 * 1 + 2 * 2 = 5

Example 2:

Input: c = 3
Output: false

Constraints:

0 <= c <= 2³¹ - 1

Test Cases

1 2	class Solution: def judgeSquareSum(self, c: int) -> bool:

solution_test.py

import pytest

from solution import Solution


@pytest.mark.parametrize('c, expected', [
    (5, True),
    (3, False),

    (18, True),
])
class Test:
    def test_solution(self, c, expected):
        sol = Solution()
        assert sol.judgeSquareSum(c) == expected

Thoughts

不妨设 a <= b。显然 b 最大不会超过 √c，a 最小可以是 0。这时候如果 a² + b² 刚好等于 c，说明 c 是平方和。否如根据当前平方和与 c 的大小关系，自增 a 或者自减 b 之后继续比对。

也可以直接从 a = 0 开始，每次判断 c - a² 是否是平方数，直到 a 比后者更大。

Code

solution.py

class Solution:
    def judgeSquareSum(self, c: int) -> bool:
        a = 0
        a2 = 0
        b = int(c ** 0.5)
        b2 = b * b
        while a <= b:
            if (s := a2 + b2) == c:
                return True
            elif s > c:
                b -= 1
                b2 = b * b
            else:
                a += 1
                a2 = a * a

        return False

本文除了题目描述（Problem）部分之外，其他内容由 Calf 原创，采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。