数列 $a_1, a_2, ... a_N$ があります。
$1\leq a_i\leq D, a_i\in \mathbb{Z} (1\leq i\leq N)$ として、${1,2,...,N}$の部分集合 $S$ であって、$\sum_{i\in S}a_i = K$ になるようなものは存在するか判定してください。
制約:
入力は整数
$1\leq N, D\leq 5000$
$1\leq K \leq N\times D$
ただの部分和問題ですが、ナップザック問題みたいな感じで解こうとすると計算量が $\mathcal{O}(N^2D)$ で大変時間がかかります。bitset高速化を用いても時間がかかりそうです。
ところが、この問題はある方法を使うと $\mathcal{O}(ND)$ で解けるらしいです。
Pisinger の方法
ここになんか載ってます。元論文の pdf は読めないようです。
貪欲パートと DP パートがあって、順に行っていきます。
貪欲パート
$i$ の小さいほうから $a_i$ を足し上げていって、和が $K-D$ 以上になったらやめます。このとき、どこまでの要素が和に組み込まれたかを保持しておきましょう。この値は後で使います。今回は、和が $\sum_{i=0}^{b-1}a_i$ になるような $b$ としておきましょう。
def Pisinger(N, D, a, K):
# b: どこまで足したか +1
# bsm: 0 から b-1 まで a[i] を足したときの和
b = 0
bsm = 0
for x in a:
if bsm < K - m:
bsm += x
b += 1
DP パート
$b$ から先を足していきます。ここでナップザック DP っぽいことをしますが、DP 配列に持つものは真偽値ではなく「部分和がこの値まで到達したか、到達したならば貪欲パートで足した要素が部分和の構成要素として残っていることが保証されるのは先頭からどの要素までか」を持ちます。
言葉で書くと意味不明なのでコードをちょっと読んでみてください。
def Pisinger(N, D, a, K):
# b: どこまで足したか
# bsm: 0 から b まで a[i] を足したときの和
b = 0
bsm = 0
for x in a:
if bsm < K - D:
bsm += x
b += 1
# 答えは存在しない
if b == N: return False
# DP パート初期化
# -1 はこの値に部分和が到達していないことを表す
# DP は [K-D, K+D] の範囲で行う
dp = [-1] * (2 * D + 1)
# bsm - (K - D) には既に部分和が到達した
# 0 から b-1 までの要素は足されたままであることが保証される
dp[bsm - (K - D)] = b
# DP
for i in range(b, N):
x = a[i]
new_dp = dp[: ]
# ふつうの DP パート
for j in range(D - 1, -1, -1):
if new_dp[j + x] < new_dp[j]: new_dp[j + x] = new_dp[j]
# 引き算の DP パート
# K より和がおおきくなってしまったもの
for j in range(2 * D, D, -1):
# 貪欲パートで足したものを引いてみる
# 足されたままのことが保証されている部分から引く
for k in range(new_dp[j] - 1, max(0, dp[j]) - 1, -1):
if new_dp[j - a[k]] < k: new_dp[j - a[k]] = k
dp = new_dp
if dp[D] == -1: return False
return True
例によって部分和が到達しないところは -1 にしています。
DP パートは二つの作業をやっています。
ふつうの DP
# ふつうの DP パート
for j in range(D - 1, -1, -1):
if new_dp[j + x] < new_dp[j]: new_dp[j + x] = new_dp[j]
dp[D] すなわち部分和が $K$ になっている部分はもう足してやる理由がないので動かしません。部分和が $K$ より大きいものにさらに足してやっても $K$ にはならないのでここも足しません。
添え字が $D$ 未満のところは部分和が $K$ 未満です。要素を足していけば $K$ になる可能性があるので足します。
普通にナップザック DP やってるだけなので計算量は $\mathcal{O}(ND)$ となります。
引き算の DP
# 引き算の DP パート
# K より和がおおきくなってしまったもの
for j in range(2 * D, D, -1):
# 貪欲パートで足したものを引いてみる
# 足されたままのことが保証されている部分から引く
for k in range(new_dp[j] - 1, max(0, dp[j]) - 1, -1):
if new_dp[j - a[k]] < k: new_dp[j - a[k]] = k
このパートでは、貪欲パートで足した要素を引いていきます。
引いていくことで和が $K$ になる可能性があるのは現在の部分和が $K$ より大きいものだけです。
まず、DP 配列に記録されている「貪欲パートで足した要素が部分和の構成要素として残っていることが保証されるのは先頭からどの要素までか」を参照します。厳密に、
$$0\leq j\lt dp[i]$$
dp[i] には上記の $j$ が和として含まれていることが保証されます。dp[i] より添え字の大きい要素も含まれるかもしれませんが、あまり気にしません。
とりあえず和の構成要素として残っているものを一つ選んで引いてみて、その結果到達した部分和では dp の値を「引いた要素の添え字の最大値」にしてあげます。
こうすると何が良いかというと、DP としては「貪欲パートで足した要素が部分和の構成要素として残っていることが保証される」＝「引くことのできる数」が最大限確保された状態を持たせることができます。
驚くべきことですが、なんか $\mathcal{O}(N^2D)$ に見えるこのパートの計算量は $\mathcal{O}(ND)$ です。DP 配列の各要素の更新回数が高々 $N$ 回で、配列の大きさが $D$ であることからわかります。
全体として
貪欲パートは $\mathcal{O}(N)$、DP パートはふつうの DP が $\mathcal{O}(ND)$、引き算の DP も $\mathcal{O}(ND)$ ですから部分和問題が計算量 $\mathcal{O}(ND)$ で解かれます。
Pisinger の方法（復元）
先のリンクにもありましたが、この方法はなんと DP の復元までやってのけます。時間的余白が少ないのでこれはまた今度です。
計算量は同じく $\mathcal{O}(ND)$ となります。
おわり
この記事のきっかけは某 ABC の G 問題です。ユーザー解説を書いてくださった cai_lw さん、まことにありがとうございます。