euler/python/e061.py

def get_four_digit_numbers(function):
    r = []
    n = 1
    f = function
    while f(n) < 10000:
        if f(n) > 999:
            r.append(f(n))
        n += 1
    return r


def search_solution(aggregator, polygonals):
    if not polygonals:
        if is_cyclic(aggregator[-1], aggregator[0]):
            return aggregator
        else:
            return []

    if not aggregator:
        for polygonal in polygonals:
            for number in polygonal:
                aggregator.append(number)
                s = search_solution(
                    aggregator, [p for p in polygonals if p != polygonal])
                if s:
                    return s
                aggregator.pop()

    for polygonal in polygonals:
        for number in polygonal:
            if is_cyclic(aggregator[-1], number) and number not in aggregator:
                aggregator.append(number)
                s = search_solution(
                    aggregator, [p for p in polygonals if p != polygonal])
                if s:
                    return s
                aggregator.pop()
    return []


def is_cyclic(a, b):
    return str(a)[-2:] == str(b)[:2]


def euler_061():
    triangles = get_four_digit_numbers(lambda n: n * (n + 1) // 2)
    squares = get_four_digit_numbers(lambda n: n**2)
    pentas = get_four_digit_numbers(lambda n: n * (3 * n - 1) // 2)
    hexas = get_four_digit_numbers(lambda n: n * (2 * n - 1))
    heptas = get_four_digit_numbers(lambda n: n * (5 * n - 3) // 2)
    octas = get_four_digit_numbers(lambda n: n * (3 * n - 2))
    s = search_solution([], [triangles, squares, pentas, hexas, heptas, octas])
    s = sum(s)
    return s


if __name__ == "__main__":
    print("e061.py: " + str(euler_061()))
    assert(euler_061() == 28684)