LeetCode 633. Sum of Square Numbers

September 13, 2020

Naively, we can do this in $O(n^2)$ by iterating through all pairs of integers, squaring them, and comparing them to $c$ .

Or, we can go faster by iterating through integers, $i$ , and checking whether ( $c$ - $i^2$ ) has an integer square root. The largest possibility we would look at is $sqrt(c/2)$ , and with the $sqrt$ function taking $O(log(n))$ time, this improvement would run in $O(sqrt(n)log(n))$ .

There might be a mathematical trick for a faster solution, but I can't think of it.

import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertEqual(self.s.judgeSquareSum(5), True)
        self.assertEqual(self.s.judgeSquareSum(3), False)

    def test_small_numbers(self):
        self.assertEqual(self.s.judgeSquareSum(0), True)
        self.assertEqual(self.s.judgeSquareSum(1), True)
        self.assertEqual(self.s.judgeSquareSum(2), True)
    
    def test_big_numbers(self):
        self.assertEqual(self.s.judgeSquareSum(48100), True)
        self.assertEqual(self.s.judgeSquareSum(48101), False)

from math import isqrt
class Solution:

    def judgeSquareSum(self, c: int) -> bool:
        for i in range(isqrt(c)+1):
            p = i ** 2
            if isqrt(c - p) ** 2 + p == c:
                return True
        return False

if __name__ == "__main__":
    unittest.main()