ProjectEuler-061

#!/usr/bin/python
#-*-coding:utf-8-*-

#problem 61

from itertools import permutations

F = (lambda x: x*(x+1)/2, lambda x: x**2, lambda x: x*(3*x-1)/2, lambda x: x*(2*x-1), lambda x: x*(5*x-3)/2, lambda x: x*(3*x-2))
_numbers = dict(zip(range(6), [set() for n in range(6)]))
n = 1
start = 5
while start >= 0:
	for x in xrange(start, -1, -1):
		f = F[x](n)
		if f < 1000:
			break
		if f > 9999:
			start -= 1
			continue
		f = str(f)
		_numbers[x].add(f)
	n += 1

for O in permutations(range(6)):
	numbers = dict(zip(range(6), [_numbers[i] for i in O]))
	L = dict(zip(range(6), [set() for n in range(6)]))
	R = dict(zip(range(6), [set() for n in range(6)]))

	for x in numbers:
		lx, rx = (x+1) % 6, (x+5) % 6
		for s in numbers[x]:
			l, r = s[:2], s[2:]
			L[lx].add(r)
			R[rx].add(l)

	N = reduce(lambda x, y: x*y, [len(numbers[i]) for i in numbers], 1)

	while N > 1:
		for x in xrange(6):
			lx, rx = (x+1) % 6, (x+5) % 6
			L[lx] = set()
			R[rx] = set()
			X = set()
			for s in numbers[x]:
				l, r = s[:2], s[2:]
				if (l in L[x]) and (r in R[x]):
					X.add(s)
					L[lx].add(r)
					R[rx].add(l)
			numbers[x] = X.copy()
			if len(X) == 0:
				break
		n = reduce(lambda x, y: x*y, [len(numbers[i]) for i in numbers], 1)
		if n  == N:
			break
		N = n
	if N == 1:
		print "Solution found: %s => %d" % (numbers.values(), sum([sum(map(int, x)) for x in numbers.values()]))
		break