Teaching Kids Programming – Count Square Sum (Pythagorean) Triples

Teaching Kids Programming: Videos on Data Structures and Algorithms

A square triple (a,b,c) is a triple where a, b, and c are integers and a^2 + b^2 = c^2. Given an integer n, return the number of square triples such that 1 <= a, b, c <= n.

Example 1:
Input: n = 5
Output: 2
Explanation: The square triples are (3,4,5) and (4,3,5).

Example 2:
Input: n = 10
Output: 4
Explanation: The square triples are (3,4,5), (4,3,5), (6,8,10), and (8,6,10).

Constraints:
1 <= n <= 250

Count Square Sum Triples via Quadratic Bruteforce Algorithm

The most intutive bruteforce would be O(N^3) to iterate all possible (a,b,c) triplets to see if they form relations:

class Solution:
    def countTriples(self, n: int) -> int:
        ans = 0
        for i in range(1, n+1):
            for j in range(1,n+1):
                for k in range(1,n+1):
                    if i**2+j**2==k**2:
                        ans+=1
        return ans

Square Sum Triplets aka. Pythagorean Triplets: where a, b and c are positive integers. We can bruteforce pair (a, b) and check if c is integer and less or equal than n. We can use the sqrt function or binary search to find out if c is a square number. The time complexity is O(N^2) quadratic considering the sqrt is O(1).

class Solution:
    def countTriples(self, n: int) -> int:
        ans = 0
        for i in range(1, n + 1):
            for j in range(1, n + 1):
                k = i * i + j * j
                c = int(k ** 0.5)
                if c * c == k and c <= n:
                    ans += 1
        return ans

a can’t be equal to b, because it is is, i.e. therefore c is not a integer. We can improve the runtime to 50% by only iterating the pairs (a, b) where a is smaller than b – and increment the answer by two if we find a integer square number c – (a, b, c) and (b, a, c).

class Solution:
    def countTriples(self, n: int) -> int:
        ans = 0
        for i in range(1, n + 1):
            for j in range(1, i):
                k = i * i + j * j
                c = int(k ** 0.5)
                if c * c == k and c <= n:
                    ans += 2
        return ans

Alternatively, we can pre-compute the square numbers and mark them in an array, then we can iterate the (a, b) pair and check if a*a+b*b is a square number by checking in the pre-computed boolean array.

class Solution:
    def countTriples(self, n: int) -> int:
        ans = 0
        arr = [False] * (n * n + 1)
        for i in range(n + 1):
            arr[i * i] = True
        for i in range(1, n + 1):
            j = i
            while j * j + i * i <= n * n:
                if arr[i * i + j * j]:
                    ans += 2
                j += 1
        return ans

The time complexity is O(N^2) and the space complexity is O(N^2) but this method doesn’t require sqrt calculation and may thus be a bit faster.

–EOF (The Ultimate Computing & Technology Blog) —