部分和(bitset)
和 $x$ を作れるかを整数の $x$ ビット目で表し、要素 $a_i$ を加えるたびに bits |= bits << a_i とします。シフトと OR がまとめて行われるので、普通の DP よりかなり速くなります。
subset_sums(a, s): 非負整数の列 $a$ からいくつかを選んだ和として作れる $s$ 以下の値を、ビットの集合(整数)で返す。和 $x$ を作れるのは(res >> x) & 1が $1$ のとき。 $O(Ns / w)$subset_sum_restore(a, s): 和がちょうど $s$ になる選び方を 1 つ、添字の昇順のリストで返す。作れなければNone。途中の結果をすべて持つので、メモリは $Ns$ ビットほど使う。 $O(Ns / w)$
ここで $w$ はワードの大きさ( $64$)です。
例( $a$ を 2 つに分けたときの、和の差の最小値):
total = sum(a)
bits = subset_sums(a, total // 2)
x = bits.bit_length() - 1 # total // 2 以下で作れる最大の和
ans = total - 2 * x
コード
def subset_sums(a, s):
# a のいくつかを選んだ和として作れる s 以下の値を、ビットの集合(整数)で返す
# 和 x を作れる <=> (res >> x) & 1
mask = (1 << (s + 1)) - 1
bits = 1
for x in a:
bits = (bits | bits << x) & mask
return bits
def subset_sum_restore(a, s):
# 和がちょうど s になる部分集合を、添字のリストで 1 つ返す。作れなければ None
mask = (1 << (s + 1)) - 1
dp = [1]
for x in a:
dp.append((dp[-1] | dp[-1] << x) & mask)
if not dp[-1] >> s & 1:
return None
res = []
for i in range(len(a) - 1, -1, -1):
if not dp[i] >> s & 1:
res.append(i)
s -= a[i]
res.reverse()
return res