各桁が異なる自然数について（3）

各桁が異なる自然数について（2）の続きです。

すべての東工大数の平方の和を考えましょう。
平方なので各桁ごとに和を取るというわけにもいきません。ここでは分割統治法を使ってみます。例えば1〜5を1回ずつ使った5桁の数の平方の総和を考えます。120個ありますね。5桁を2つに2桁と3桁にわけます。そして上の2桁は1, 2、下の3桁は3, 4, 5を使うとします。この分け方は10通りあります。これを単純に計算すると、

12345² = 12² * 10⁶ + 2 * 12 * 345 * 10³ + 345²
12354² = 12² * 10⁶ + 2 * 12 * 354 * 10³ + 354²
…
21543² = 21² * 10⁶ + 2 * 21 * 543 * 10³ + 543²

これを足し合わせます。そうすると、12² * 10⁶など右辺の第1項は6回、345²などの第3項は2回出てきます。それから第2項を全て足すと、

2 * (12 + 21) * (345 + 354 + 435 + 453 + 534 + 543) * 10³

となります。各和は前回の方法で直ちに求められます。
zero-leadingの問題があるのでもうちょっと複雑ですが、基本的にはこの考え方で、あとメモ化を使うと、単に列挙して平方したときの100倍くらい速かったです。

from itertools import *
import time

def factorial(n):
	return reduce(lambda x, y: x * y, xrange(1, n + 1), 1)

def sum_perms(ds, head):
	L = len(ds)
	if L == 1:
		return ds[0]
	elif head and 0 in ds:
		return sum_perms(ds, False) - sum_perms(ds[1:], False)
	else:
		return factorial(L) * (10 ** L - 1) / 9 * sum(ds) / L

memo = { }
def sum_square_perms(ds, head):
	L = len(ds)
	if L == 1:
		return ds[0] ** 2
	
	head = head and 0 in ds
	if (ds, head) in memo:
		return memo[(ds,head)]
	
	L1 = L / 2
	L2 = L - L1
	s1, s2, s3, s4 = 0, 0, 0, 0
	for ds1 in combinations(ds, L1):
		ds2 = tuple(d for d in ds if d not in ds1)
		s1 += sum_square_perms(ds1, head)
		s2 += sum_perms(ds1, head) * sum_perms(ds2, False)
		if head and 0 in ds1:
			s3 += sum_square_perms(ds2, False)
		else:
			s4 += sum_square_perms(ds2, False)
	s = (s1 * 100 ** L2 * factorial(L2) +
		 s2 * 2 * 10 ** L2 +
		 s3 * factorial(L1 - 1) * (L1 - 1) +
		 s4 * factorial(L1)) % D
	memo[(ds,head)] = s
	return s

t0 = time.clock()
D = 1000000007
print sum(sum_square_perms(ds, True) for n in tuple(xrange(1, 11))
								for ds in combinations(range(10), n)) % D
print time.clock() - t0