hip_checker.py

#!/usr/bin/env python3
#
# Solution to the 538 Riddler Express for 2020-10-02
# https://fivethirtyeight.com/features/can-you-eat-all-the-chocolates/

from math import sqrt

# The number we're checking for "hippness"
target = 1400

# Create a set of perfect squares from 1 to the point where successive
# squares are > the target apart. Since the distance between the Nth square
# and the (N+1)st square is 2*n+1, we can just do first target/2+1 squares
squares = set([i**2 for i in range(target//2 + 2)])

# Now for each square, see if (target+square) is also a perfect square (thus,
# target+square-square == target). This is faster than checking each pair of
# squares 
n_found = 0
for s in squares:
    if s + target in squares:
        n_found += 1
        print("{} - {} = {} = {}**2 - {}**2".format(s+target, s, target,
                                                    int(sqrt(s+target)),
                                                    int(sqrt(s))))
print("{} pair(s) found".format(n_found))