# Pentagonal numbers are generated by the formula, Pn=n(3n−1)/2. The first ten pentagonal numbers are:
#
# 1, 5, 12, 22, 35, 51, 70, 92, 117, 145, ...
#
# It can be seen that P4 + P7 = 22 + 70 = 92 = P8. However, their difference, 70 − 22 = 48, is not pentagonal.
#
# Find the pair of pentagonal numbers, Pj and Pk, for which their sum and difference are pentagonal and
# D = |Pk − Pj| is minimised; what is the value of D?


from math import sqrt


def is_pentagonal(n):
    if n < 1:
        return 0
    k = (1 + sqrt(1 + 24 * n)) / 6
    if k == int(k):
        return int(k)
    return 0


pentagonals = dict()


def pentagonal(n):
    global pentagonals
    if n in pentagonals:
        return pentagonals[n]
    p = int(n * (3 * n - 1) / 2)
    pentagonals[n] = p
    return p


def main():
    a = 1
    results = []
    while a < 3000:
        # Ugly empirical upper bound
        # No guarantee that the difference is indeed minimal but the solution is correct, so oh well
        pa = pentagonal(a)
        b = 1
        while b < a:
            pb = pentagonal(b)
            if is_pentagonal(pa - pb) and is_pentagonal(pa + pb):
                print(b, a, pa - pb)
                results.append([b, a, pa - pb, pb, pa, pa + pb])
            b += 1
        a += 1
    return results


print(main())