CodeIQ:「プラス・マイナス・ゼロ」問題

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「「プラス・マイナス・ゼロ」問題(CodeIQ)を参照してください。

問題の概要

問題を引用します。
自然数 n に対して、次の等式を考えます。
□1□2□3□4…□n = 0
四角の部分には、プラス(+)またはマイナス(-)の記号が入ります。
等式を成立させるような記号の入れ方を考えましょう。

例えば、n = 7 のとき、次の 8 通りの入れ方が可能です。
+1+2-3+4-5-6+7 = 0
+1+2-3-4+5+6-7 = 0
+1-2+3+4-5+6-7 = 0
+1-2-3-4-5+6+7 = 0
-1+2+3+4+5-6-7 = 0
-1+2-3-4+5-6+7 = 0
-1-2+3+4-5-6+7 = 0
-1-2+3-4+5+6-7 = 0

自然数 n に対し、等式を成立させる記号の入れ方の数を F(n) と定義します。
例えば F(7) = 8 です。
同様に、F(3) = 2、F(8) = 14 となることが確かめられます。

標準入力から、自然数 n (1 ≦ n ≦ 50)が与えられます。
標準出力に F(n) の値を出力するプログラムを書いてください。

私のプログラム

Pythonで解答しています。

#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

#==============================================================================
'''
	自然数 n に対して、次の等式を考えます。
		□1□2□3□4…□n = 0
	四角の部分には、プラス(+)またはマイナス(-)の記号が入ります。
	等式を成立させるような記号の入れ方を考えましょう。
'''
#------------------------------------------------------------------------------
import fileinput
import math

# nを(m,n-m)に分割した時、mによって何通りに分割可能かを記録する
# 次のようになる
#	      m
#         0  1  2  3  4  5  6  7
#	n 0   1
#	  1   1  0
#	  2   1  0  0
#	  3   1  1  0  0
#	  4   1  1  0  0  0
#	  5   1  1  1  0  0  0
#	  6   1  1  1  1  0  0  0
#	  7   1  1  1  2  0  0  0  0
Memo = []
Memo.append([1])	# n=0の場合

#	partInt(mx)
#	mx 元の数値
#	base base以下の値に分割
def partInt(mx, base):

	for n in range(1,mx+1):
		# mxを2つの数に分割した場合に必要な領域を作成する
		# 無駄があるがとりあえず気にしない
		# mx=3の場合、[3,0],[2,1][1,2][0,3]
		Memo.append([0 for i in range(n+1)])

		# 分割の上限を決定する(OEIS A008289のPROGのロジックから)
		# 参考のロジックとはループ方向が逆になるのでxを求める式が変わる
		x = math.ceil((math.sqrt(8*n+1)-1)/2)-1
		lim = n - x

		# 問題のため計算開始位置を決定する
		a = base
		if lim < base:
			base = lim-1

		for i in range(base, lim):
			# Memo[n][i]の組み合わせ数を求める
			# Memo[i]の合計値を求める開始位置を決める
			# n-i <= iの場合?以降の数値を合計する
			# 例えばn=6で(3,3)に分割した場合、i=3の分割数のうち[3,0]は不要[2,1]は必要
			st=0
			if i - (n-i) >= 0:
				st = 2*i-n+1

			Memo[n][i] += sum(Memo[i][st:])

# 答えを計算する
# h 入力値nに関して[1..n]の合計値の1/2の値
# ipt 入力値
def getAns(h, ipt):
	tgt = Memo[len(Memo)-1]
	s = sum(tgt[h-ipt+1:])
	return s*2

if __name__ == "__main__":
	for line in fileinput.input():
		if not line.strip():
			continue

		n = int(line.strip())
		s = sum(list(range(1,n+1)))
		h=0
		if s%2:
			print(0)
		else:
			h = s//2
			partInt(h,n)
			print(getAns(h, n))

解説

巨大な組み合わせを求める問題です。単純に考えると与えられた数nに対し1〜(n-1)個の数を選んだ場合、選んだ数と残りの数の合計が等しくなる組み合わせを求めるということですが膨大な計算量になってしまいます。

基本的な考え方

まず、選んだ数と残りの数が等しくなるというのは選んだ数の合計が1〜nの合計の1/2に等しいということと同じです。問題の例(n=7)の場合1〜7の合計は28でその半分は14です。例のプラスがついた数字だけを拾うと14になります。

次に1〜nの合計の半分をhとすると、この問題でhを構成する数は、「hをn以下の互いに異なる数で分割したもの」と同義です。分割を効率良く求められれば問題を解けます。

main

mainの76行目で1〜nの合計を求めています。78行目で合計が2で割り切れない場合を弾いています。合計が2で割り切れないということは±0になる組み合わせが存在しないことなので計算の必要がありません。
2で割り切れた場合はpartInt()に合計値の半分とnを渡して分割を行います。
分割が終わったらgetAns()で答えを求めて印字します。

partInt()

この関数で自然数分割を行います。計算を高速にするため(多分)動的計画法を使っています。
グローバル変数Memoはコメントに書いてある通り、0〜合計値の半分までの全ての数を互いに異なる数に分割した場合、何パターンに分割できるのかを記録します。2次元配列になっていてnの行はnを(n,0)、(n-1,1)、(n-2,2)…(n-m,m)…(0,n)に分割し、mを更に分割した場合、それぞれのmについて何通りの分割があるかを保持します。従って、mが今回必要な分割数の上限に相当します。

処理の手順ですが、ややこしいので具体的に示します。partInt(8,4)の場合を考えます。
まず、1〜8までの数を分割します。欲しい結果は8の場合だけですが、例えば8を(5,3)に分割した時に3は更に分割できます。この時に3の分割を計算せず計算済みの値を利用するため1から順に分割を行い、結果を利用できるようにメモします。この1〜8までの繰り返しが25行目からのループです。

39行目でメモの領域を作成しています。コメントに書いてあるだけの領域を作成しますが、互いに異なる数に分割するので不要な領域があります。コメントのように3を分割する場合、[(3,0), (2,1)]の分割は必要ですが[(1,2),(0,3)]の分割は必要ありません。無駄ですが大したことはないので作成しますが無視します。

43〜49行目はどこまで分割するのかの上限を決めています。前述の3を分割した場合の例で言うと分割は[(3,0), (2,1)、(1,2),(0,3)]とできますが[(1,2),(0,3)]は不要なので、[(3,0), (2,1)]までで51行目からのループを終了するために計算しています。これはOEISのA008289を参考にしました。

51行目からの処理はある数を2つ(n-m, m)に分けた時にmによっていくつ組み合わせがあるかを計算しています。この時により少ない値の分割結果が役に立ちます。例えば7を分割した場合、[(7,0), (6,1), (5,2), (4,3)]に分割します(※(2,5)以降は43〜49行目の上限設定により計算されません)。これはMemo[7][0]〜Memo[7][3]に相当します。
この時、Memo[0]〜Memo[6]はすでに分割が計算されています。

0は初期値でMemo[0]=[0]に設定されているのでMemo[7][0]=1です。
過去の計算結果でMemo[1]=[1]になっているのでMemo[7][1]=1です。
同様にMemo[2]=[1]なのでMemo[7][2]も1パターンです。
Memo[3]=[1,1]です。3は(3,0)と(2,1)に分割できるのでMemo[3][0]=1、Memo[3][1]=1だからです。なのでMemo[7][3]=sum(Memo[3])=2になります。
このように計算することで数を2つに分割するだけで分割のパターン数を計算することができます。

getAns()

答えを計算すします。
MemoはMemo[入力値の合計の1/2][]の2次元配列なので最後の行に答えが含まれます。そしてMemo[h][0],Memo[h][1],Memo[h][2],…,Memo[h][n],…はn以降がnを最大値として分割した場合のパターン数になります。
私のプログラムではn+1以降の合計の2倍を返していますが、問題の例で言うと最大6で分割した数で、これは6がプラス側だけを選択したのと同じです。そのパターンをプラス・マイナス反転したのが全体のパターン数になります。
わざわざこんなことをしましたが、よく考えるとn以降を合計した場合とn+1以降の合計の2倍は同じ結果になります。

雑感

実はこの問題と同じ頃に別の問題でも自然数分割(そっちはnをm個以下の互いに異なる自然数に分割する)を使っていて、そちらを先にやっていたのでロジックを流用できそうでしたができず、今回のロジックは新規に考え直しました。
私としてはこの問題は非常にむずかしかったと思います。自分で考えたコードでこれだけコメントを入れているにもかかわらず見直したら、これは何をやってたんだっけ? と考え直さなければならないところがあります。