Project Euler & HackerRank problem 44 solution: Find the pair of pentagonal numbers, Pj and Pk, for which their sum and difference is pentagonal and |Pk − Pj| is minimised - solved using Python

Project-Euler-Problem-44.py

#HackerRank Project Euler 44 version
from math import sqrt
isp = lambda p: (sqrt(24*p+1)+1)/6 % 1 ==0

ps = [(3*i*i - i) // 2 for i in range(1, 1000005)]
n,k = map(int, input().split())
for n in range(k+1, n+1):
    if isp(ps[n]+ps[n-k]) or isp(ps[n]-ps[n-k]): print (ps[n])