This problem just required a bit of algebra. My strategy was to just go
through all possible values of i, j, k for which i**2 - j**2 - k**2 was
greater than zero and less than a million. The first thing I did was remove k
from that equation because we know that k = 2*j - i. The next thing was to
figure out the upper bound for j. We know that my value for the variable key
needs to be greater than zero. From there it is easy to solve for j in terms
of i. Next I checked the output of i and j to see when the values were in the
correct range. That allowed me to set an upper bound for i, and know that at a
certain point, we can break out of j. So, in the end I was able to find the
lower and upper bounds for both i and j, with the addition of breaking out
early from j when k falls too low. Runs in under a second.
def solve(n, m):
solutions = {}
for i in range(3, 5*m//4+1):
for j in range(4*i//5, i//2, -1):
key = 4*i*j - 5*j**2
if key >= m:
break
solutions[key] = solutions.get(key, 0) + 1
count = 0
for val in solutions.values():
count += val == n
return count
if __name__ == '__main__':
import sys
n = eval(sys.argv[1])
m = eval(sys.argv[2])
print(solve(n, m))
TESTS:
import pytest
from problem import solve
_test_solve = (
(1, 100, 25),
(2, 100, 8),
(3, 100, 8),
(4, 100, 2),
(5, 100, 0),
(1, 200, 44),
(2, 200, 20),
(3, 200, 16),
(4, 200, 5),
(5, 200, 2),
(6, 200, 0),
)
@pytest.mark.parametrize('n,m,expect', _test_solve)
def test_solve(n, m, expect):
assert expect == solve(n, m)