超軽実装！個数制限付き部分和問題 (数え上げ) を 10 行未満で実装する

## はじめに

本記事は，タイトルの通り__個数制限付き部分和問題 (数え上げ)__を超軽実装で解くことを目的としています．

## 部分和問題とは

これから扱う問題の概要を以下に示しておきます．

__【問題概要】__

正整数列 $A_0,\dots,A_{N-1}$ と正整数 $S$ が与えられた下で， $$\sum\_{i=0}^{N-1} c\_i\cdot A\_i=S$$ を満たす非負整数の組 $(c_0,\dots,c_{N-1})$ の個数を求めよ．

------

題中に現れた $c_i$ には制約が加えられることが多く，特に，以下に示す三種類は頻出です．

CASE 1 : $c_i\in\\{0,1\\}$
CASE 2 : $0\leq c_i$
CASE 3 : $0\leq c_i\leq M_i$

そして，この記事で扱うのは 3 つ目のケースです．つまり，$A_i$ を使って良い個数に上限が定められているケースです．問題の難しさは一般的には 1<2<3 だと思うので，1，2 の解法についても簡単に説明しておきます．

### CASE 1 : 選ぶ個数は 0 個または 1 個

この場合は典型的な動的計画法により解くことが出来ます．

$1\leq i\leq N$，$0\leq j\leq S$ に対して，DP テーブルを

$$dp[i][j]=A_{i-1} \text{ まで見て，和が } j \text{ であるような場合の数 }$$

と定義し，$i=0$ の時は

$$
dp[0][j]=\begin{cases}
1& \text{(if}\ \ j=0 \text{)}\\\ 
0& \text{(otherwise)}
\end{cases}
$$

とします ( 初期状態で和は $0$ なので，$j=0$ では $1$ 通り，それ以外では $0$ 通りだと考えると自然だと思います )．

すると，この DP テーブルは $i=0,\dots,N-1$ に対して以下のように更新すればよいことが分かります．

$$
dp[i+1][j]=\begin{cases}
dp[i][j]&\text{(if}\ \ j\lt A_i\text{)} \\\
dp[i][j]+dp[i][j-A_i]&\text{(otherwise)}
\end{cases}
$$

最終的に求める答えは $dp[N][S]$ となり，全体の計算量は $O(NS)$ です．以下に Python でのコードを掲載します．

```python
# 入力部は省略
DP = [[0] * (S + 1) for i in range(N + 1)]
DP[0][0] = 1
for i in range(N):
    for j in range(S + 1):
        DP[i + 1][j] = DP[i][j] + (j >= A[i] and DP[i][j - A[i]])
print(DP[N][S])
```

更に，DP テーブルの更新を in-place にすることで場合分けを無くすことができるので実装は簡単になります (メモリ量削減という意味の方が本質的ではありますが)．しかし，この場合は注意が必要で，`j` のイテレーションの方向を降順にする必要があります．コードは以下のようになります．

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(S, A[i] - 1, -1):
        DP[j] += DP[j - A[i]]
print(DP[S])
```

余談ですが，python は負の index が使えるのでこんな書き方も出来ます (コードゴルフ的ですね)．

```python
DP = [1] + [0] * S
for v in A:
    for j in range(S - v + 1):
        DP[~j] += DP[~j - v]
print(DP[S])
```

### CASE 2 : 選ぶ個数は 0 個以上 (上限無し)

同様の DP テーブルの定義を考えると，以下のように更新できることが分かります．

$$dp[i+1][j]=\sum_{k=0}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]$$

しかし，このまま計算すると $\forall i\in \\{0,\dots,N-1\\}. A_i=1$ のようなケースで $O(NS^2)$ かかるので計算量が悪化しています．

そこで，少し式変形をします．$\lfloor j/A_i\rfloor\geq 0$ に注意します．

$$
\begin{aligned}
dp[i+1][j]
&=\sum_{k=0}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]\\\
&=dp[i][j]+\sum_{k=1}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{\lfloor j/A_i\rfloor-1}dp[i][j-(k+1)\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{\lfloor (j-A_i)/A_i\rfloor}dp[i][(j-A_i)-k\cdot A_i]\\\
&=\begin{cases}
dp[i][j]& \text{(if}\ \ j<A_i\text{)}\\\
dp[i][j]+dp[i+1][j-A_i]& \text{(otherwise)}
\end{cases}
\end{aligned}
$$

選ぶ個数が $0$ 個または $1$ 個のケースと似た式を得ました．この式をよく見ると，$\mathrm{mod}\ {A_i}$ 毎に累積和を取っていることが分かります．

この更新式を用いることで，テーブル全体を $O(NS)$ で更新することができ，CASE 1 と比べて計算量を悪化させることばく答えを求めることが出来ました．以下に Python でのコードを示します．

```python
# 入力部は省略
DP = [[0] * (S + 1) for i in range(N + 1)]
DP[0][0] = 1
for i in range(N):
    for j in range(S + 1):
        DP[i + 1][j] = DP[i][j] + (j >= A[i] and DP[i + 1][j - A[i]])
print(DP[N][S])
```

このケースでも，in-place な更新をすることで実装を容易にし，メモリ量も削減することが出来ます．CASE 1 では `j` の降順に更新していましたが，今度は昇順に更新するようにすればよいです．

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(A[i], S + 1):
        DP[j] += DP[j - A[i]]
print(DP[S])
```

あるいはこうでしょうか．

```python
DP = [1] + [0] * S
for v in A:
    for j in range(S - v + 1):
        DP[j + v] += DP[j]
print(DP[S])
```

これで CASE 1 と CASE 2 の場合は解けました．次節ではいよいよ本題の CASE 3 (個数制限付き部分和問題) を考えます．

## 個数制限付き部分和問題 (数え上げ)

ここでは，$A_i$ を選べる個数の上限値 $M_i\geq 0$ が定められている場合を考えます．

個数が無制限の場合は，DP テーブルを

$$dp[i][j]=A_{i-1} \text{ まで見て，和が } j \text{ であるような場合の数 }$$

で定義すると，計算量は悪いですが，

$$
dp[i+1][j]=\sum_{k=0}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]
$$

で更新できました．個数制限がある場合は，これを少し変形して

$$
dp[i+1][j]=\sum_{k=0}^{\min\\{\lfloor j/A_i\rfloor,M_i\\}}dp[i][j-k\cdot A_i]
$$

とすればよいです．個数制限無しのケースで述べましたが，基本的には $\mathrm{mod}\ A_i$ での累積和のような見た目をしています．

式中に $\min$ が入っていて少しわかりづらいので，今回は図を使います．

以下は $S=32,\ A_i=5,\ M_i=2,\ j=7, 13, 19, 20, 31$ での例です．

![subset_sum.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/618590/4d46a1ef-cfa4-acbd-4f5f-30d3b3bd276d.png)

DP テーブルの各要素は $\mathrm{mod}\ A_i$ 毎に分けた列における区間和となっています．そして，赤セルの添字 $j$ と，$j$ と同列にある青セルの最大の添字 $k$ の間には，$k=j-(M_i+1)\cdot A_i$ という関係が成り立っていることが分かります．従って，一度累積和を取ったあと，$j$ の降順に $j-(M_i+1)\cdot A_i$ までの累積和を引いていけばよいです．

以上より，コードは以下のようになります (in-place な実装です)．

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(A[i], S + 1):
        DP[j] += DP[j - A[i]]
    d = (M[i] + 1) * A[i]
    for j in range(S, d - 1, -1):
        DP[j] -= DP[j - d]
print(DP[S])
```

複雑そうな問題ですが，かなり簡潔に実装できていると思います．

また，かなり蛇足なんですが，軽実装がテーマということで短くしてみました (コードゴルフはちゃんとやったことが無いので，強い人に怒られそうですが)．変数名まで変えると完全に意味が分からないので `DP` はそのままですが，129 byte です．

```python
DP=[1]+[0]*S
for a,m in zip(A,M):
 for j in range(-a-~S):DP[j+a]+=DP[j]
 for j in range(~m*a-~S):DP[~j]-=DP[~j+~m*a]
print(DP[S])
```

## 最後に

ここまで読んで頂き，ありがとうございました．紹介した実装は意外と知られてないのかな，と思い執筆しましたが，使い古されたものであれば申し訳ありません．

また，記事中では触れませんでしたが，形式的冪級数という道具を用いれば今回紹介した DP 更新式がかなり容易に求まります．FFT と組み合わせれば計算量的にも強力な道具なので，興味があれば勉強するとよさそうです (私はまだほとんど何も分かっていません)．

終わりに，個数が $0$ 個または $1$ 個の場合のコード，個数無制限のコード，個数制限付きのコードを再掲しておきます．

- 個数が $0$ 個または $1$ 個

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(S, A[i] - 1, -1):
        DP[j] += DP[j - A[i]]
print(DP[S])
```

- 個数無制限

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(A[i], S + 1):
        DP[j] += DP[j - A[i]]
print(DP[S])
```

- 個数制限付き

```python
DP = [1] + [0] * S
for i in range(N):
    for j in range(A[i], S + 1):
        DP[j] += DP[j - A[i]]
    d = (M[i] + 1) * A[i]
    for j in range(S, d - 1, -1):
        DP[j] -= DP[j - d]
print(DP[S])
```

本当の終わりに，個数制限付きの場合での DP 更新式の式による導出を載せておきます．

<details><summary>(折りたたみ)</summary><div>

$\lfloor j/A_i\rfloor$ と $M_i$ の大小関係で場合分けをします．

- $\lfloor j/A_i\rfloor \leq M_i$ の場合

$\lfloor(j-A_i)/A_i\rfloor\leq\lfloor j/A_i\rfloor\leq M_i$ なので，個数無制限の場合と同様の式変形が出来ます．

$$
\begin{aligned}
dp[i+1][j]
&=\sum_{k=0}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]\\\
&=dp[i][j]+\sum_{k=1}^{\lfloor j/A_i\rfloor}dp[i][j-k\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{\lfloor j/A_i\rfloor-1}dp[i][j-(k+1)\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{\lfloor (j-A_i)/A_i\rfloor}dp[i][(j-A_i)-k\cdot A_i]\\\
&=\begin{cases}
dp[i][j]& \text{(if}\ \ j<A_i\text{)}\\\
dp[i][j]+dp[i+1][j-A_i]& \text{(otherwise)}
\end{cases}
\end{aligned}
$$

- $\lfloor j/A_i\rfloor \gt M_i$ の場合

$\lfloor j/A_i\rfloor\gt M_i$ なので，$\lfloor j/A_i\rfloor\geq M_i+1$ です．

従って，$\lfloor (j-A_i)/A_i\rfloor\geq M_i$ で， $\min\\{\lfloor (j-A_i)/A_i\rfloor,M_i\\}=M_i$ となることに注意して式変形をします．

$$
\begin{aligned}
dp[i+1][j]
&=\sum_{k=0}^{M_i}dp[i][j-k*A_i]\\\
&=dp[i][j]+\sum_{k=1}^{M_i}dp[i][j-k\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{M_i-1}dp[i][j-(k+1)\cdot A_i]\\\
&=dp[i][j]+\sum_{k=0}^{M_i-1}dp[i][(j-A_i)-k\cdot A_i]\\\
&=dp[i][j]+dp[i+1][j-A_i]-dp[i][j-(M_i+1)\cdot A_i]
\end{aligned}
$$

以上より，テーブル全体を $O(NS)$ で更新できます．

更新式を眺めると，図から求めたように，$dp[i+1][j]$ の値は $dp[i][j]$ の $\mathrm{mod}\ A_i$ 毎の累積和から $dp[i][j-(M_i+1)\cdot A_i],dp[i][j-(M_i+2)\cdot A_i],\dots$ を引いた値になっていることが分かります．</div></details>