# 関数適用による状態遷移で DP を解く
動的計画法 (以下 DP) の実装にあたっては、一次元もしくは多次元配列を用いて、ある局面での計算結果を配列の要素にメモ化しながら計算を進めていく··· ということが多くの教科書、解説に書いている。
例えば典型的なナップザック問題であれば二次元配列を用意して、for の二重ループで二次元配列を更新していく実装になる。事前に計算された計算を足場に次の計算を進める、累積的な計算過程になっていてこれが漸化式と対応するというのが、私が最初に獲得したメンタルモデルだった。
一方、現在はこのメンタルモデルを更新して DP の計算過程を「関数適用による状態遷移」と「複数の状態の結合演算」と捉えていて、そのメンタルモデルをトレースしたような実装で解いている。
すべての DP がこの実装で解けるわけではないが、ナップザック問題と同様の計算構造のタイプの DP であればほぼ解ける。また実装云々以上に、DP を状態遷移でとらえるモデルの方が、自分としては捉えやすかった。
以下そのあらましについて記す。
## 手続き的なナップザック問題の解法
- [ナップサック問題 | アルゴ式](https://algo-method.com/tasks/342)
ナップザック問題は二次元表を元にした動的計画法で解く、というのを最初に学んだ。以下のようなもの。
例えば $(重さ, 価値)$ で表される荷物が `(5, 10), (2, 8), (3, 4)` と $3$ つあったとする。
袋にいられる荷物の最大が $5$ に制限されているとき、どの荷物を袋に入れると「袋に入った荷物の価値の総和」を最大化出来るか。
二次元表に対し
- 列を袋に入っている荷物の重さ
- 行は荷物 (のインデックス)
をそれぞれ割り当てて、表の値に袋に入っている価値の合計値を格納する。
ある荷物があったとき、その荷物に対してとれる行動は「選ぶか、選ばないか」の二択の分岐になる。
ある荷物を選ばなかったときは当然重さは変わらず、選んだときは重さが増えて、価値を足すことができる。
![[スクリーンショット 2024-02-09 14.36.02.png|600]]
荷物を選んでいくと、図の最右下のマスのように、複数の方向からそのマスに値が入って来ることがある。それら複数の値は結合して一つの値にしたい。
ここでは「価値の総和の最大化」を目指しているので、二項演算の max を使い、最大値に結合する。DP の文脈ではこの複数の値を一つにまとめることを「緩和 (Relax)」と呼ぶらしい。
この二次元表のモデルを素直に実装すると、二次元配列を用意して for の二重ループでそれを書き換えていく... というような実装になる。
ChatGPT に C++ で書かせたら以下のような実装が出てきた。書き換え可能な配列を for 文を回して書き換えていく、手続き的な実装である。
```cpp
for (int j = 0; j <= W; ++j) dp[0][j] = 0;
for (int i = 0; i < N; ++i) {
for (int j = 0; j <= W; ++j) {
if (j >= w[i]) dp[i+1][j] = max(dp[i][j-w[i]] + v[i], dp[i][j]);
else dp[i+1][j] = dp[i][j];
}
}
cout << dp[N][W] << endl;
```
これが典型的なナップザック問題の解き方かなと思う。私も最初は同様の考え方で解いていた。
そして類題を解こうとしては、二次元表の構成の仕方に四苦八苦していた。二次元表の視覚的イメージを頼りに考えていたこともあり、三次元以上になるともうお手上げという感じだった。
## 関数適用による状態遷移のモデルで捉える
Haskell で DP の問題をいろいろ解いているうちに、このメンタルモデルが少しずつ変化していった。
具体的には、DP の計算過程を状態遷移のモデルでとらえるようになった。どういうことか、順を追って説明していく。
まず、先の二次元表のモデルを以下のように少し分解してみる。
![[スクリーンショット 2024-02-09 14.36.51.png|600]]
何が変わったかというと、行と行が絵的に分離されただけである。これだけだと図の違いでしかないが「全体がくっついた一つの二次元配列」というメンタルモデルからの脱出がはじめの一歩だ。
ところで DP の計算には、おおざっぱにいうと「一つ前の計算結果のみから次の計算ができる」性質がある、そこに着目してみる。上図から、一部の計算過程のみを抜き取ろう。
![[スクリーンショット 2024-02-09 14.37.24.png|600]]
荷物 `(2, 8)` を選択するかしないかを確定した局面 (行) から、次の荷物である `(3, 4)` を選択するかどうかの局面 (行) を抜き取った。
このとき `(3, 4)` を選択した/しなかった結果は一つ前の局面である `(2, 8)` の局面 (行) だけをもとに計算することができて、それより以前の局面、つまり行は必要がない。
これを 「荷物 `(2, 8)` を判断し終えた状態」から「荷物 `(3, 4)` を判断する状態」への状態遷移... と考えることはできないだろうか?
このメンタルモデルの転換が、一番重要なポイントになる。
ある行は、その時点での袋の局面、つまり状態を表している。荷物 `(2, 8)` をどうするか判断し終えた状態とみなす。
袋の状態は分岐によって複数の可能性をもっているから、その状態はスカラーではなく複数の可能性を表現できるリストや配列などのコンテナで表現すると考える。つまり「状態空間」である。その状態空間が、次の状態へ遷移する... という風に考えることができる。
さらにブレイクダウンして、状態空間全体ではなくある単一の状態から次の状態への遷移も考えてみたい。
上図のうち、袋に荷物が重さ $2$ 入っていて、その中身の価値の総和が $8$ の状態に着目する。この袋の状態に次の荷物 `(3, 4)` を入力として渡すと、この `(3, 4)` を袋に入れるか、入れないかの二つの状態に分岐して遷移する。
![[スクリーンショット 2024-02-09 14.37.49.png|600]]
これを式で書くと、以下の様に書けるだろうか。
$f ((2, 8), (3, 4)) → [(2, 8), (5, 12)]$
袋の状態 $(2, 8)$ に次の荷物 $(3, 4)$ を入力として渡し $f$ 関数により状態を遷移させると、結果、次の荷物を選ばず袋の状態がそのまま維持された $(2, 8)$ と、袋に荷物を加えた $(5, 12)$ の二つの状態に遷移する。
これを一般化すると以下となる。
$f ((w, v), (a, b)) → [(w, v), (w + a, v + b)]$
ここまでを考えると
- ナップザック問題における DP の計算というのは、ある状態空間の状態が、次の状態空間の状態へと状態遷移するとみることができる
- さらに、状態空間のなかのある要素、つまり特定の状態に着目すると、その状態が入力を得て次の状態へ遷移する、とみることができる
となっている。
状態遷移の関数にまでブレイクダウンできたら Haskell での実装がみえてくる。
```haskell
f (w, v) (a, b) = [(w, v), (w + a, v + b)]
```
この DP の計算過程を状態遷移としてとらえる考え方をうまくトレースして DP の計算そのものを Haskell の実装に落とし込めないだろうか?
## Haskell の accumArray
突然話が変わるが、Haskell の Array には `accumArray` というとても便利な関数がある。他のプログラミング言語ではあまりみない関数だが、とても強力である。
`accumArray` は、引数に与えられた複数の `(インデックス、値)` という値を、同じインデックス値を一つの値にまとめあげた配列を構築する。具体例をみる方がわかりやすい。
```haskell
let as = accumArray (+) 0 (0, 5) [(1,1), (1, 2), (2, 2), (2, 3), (5, 1), (2, 1)]
print (elems as) -- [0,3,6,0,0,1]
```
図にすると以下のようになっている。
![[スクリーンショット 2024-02-09 15.11.21.png|600]]
同じインデックスの値が、配列の一つの要素に集められるのでそれらを結合して一つの値にする必要がある。
そのための結合演算 (二項演算) を指定することができる。上記では `(+)` を指定したので同じインデックス値は和で結合される。
(尚、ユースケース的な目線でみると、上記は頻度表を作っている計算だと思うこともできる。)
`accumArray` は二項演算で同じインデックスに基づく値を結合する関数であり、ここにモノイドみ...代数的構造を感じるが本論ではないのでいったんそこは脇においておこう。
ところで、この `accumArray` の図をみると先にみた DP の状態遷移の図に構造がよく似ていることに気づく。複数の値が合流して、二項演算により結合される ... DP の緩和の計算によく似ている。
![[スクリーンショット 2024-02-09 14.37.24.png|600]]
だいぶ話が繋がってきたかと思う。
## accumArray を使い、DP の状態遷移を実装する
`accumArray` を使った実際のコードで具体的な状態遷移を実装してみる。
こんな風に書けば、ここまでのメンタルモデルを比較的トレースできているように思う。
```haskell
main :: IO ()
main = do
-- 状態遷移前の状態 (無向な状態は minBound)
let as = listArray @UArray (0, 5) [0, minBound, 8, minBound, minBound, 12]
-- 状態遷移
transition =
[ [(w, v), (w + 3, v + 4)] | (w, v) <- assocs as, v /= minBound, inRange (0, 5) (w + 3)]
-- 遷移結果をまとめあげると次の状態になる
as' = accumArray @UArray max (minBound :: Int) (0, 5) $ concat transition
print (elems as') -- [0,-9223372036854775808,8,4,-9223372036854775808,12]
```
もう少し押し進めて $f ((w, v), (a, b)) → [(w, v), (w + a, v + b)]$ という式をより直接的に実装に表現したい。
以下のように、先にもみた状態遷移関数 `f` を別途定義してそれを適用するようにすればいい。
```haskell
main :: IO ()
main = do
-- 状態遷移の定義
let f (w, v) (a, b) = [(w, v), (w + a, v + b)]
-- 状態遷移前の状態
let as = listArray @UArray (0, 5) [0, minBound, 8, minBound, minBound, 12]
-- 状態遷移 ··· 関数 f を現在の状態に適用することで次の複数の状態に遷移
transition = [f (w, v) (3, 4) | (w, v) <- assocs as, v /= minBound, inRange (0, 5) (w + 3)]
-- 遷移結果をまとめあげると次の状態になる
as' = accumArray @UArray max (minBound :: Int) (0, 5) $ concat transition
print (elems as')
```
ここまでみて分かる通り状態遷移関数 $f$ を定義し、ある状態空間の値すべてをその $f$ によって遷移させて、その遷移結果を `accumArray` でまとめあげることで次の状態空間の状態が得られた。この方法で DP の計算過程をうまく実装することができそうだ。
実際の DP の問題が解けるように整えたのが、以下の `accumArrayDP` という関数である。
ここまで見てきたモデルをそのまま実装しただけで、それ以上特別なことはやっていない。
```haskell
-- ex) accumArrayDP @UArray f max minBound (0, wx) [(0, 0)] wvs
accumArrayDP ::
( IArray a e,
Ix v,
Eq e,
Show e,
Show v,
Show (a v e),
Foldable t
) =>
((v, e) -> x -> [(v, e')]) -> -- 状態遷移関数 f v / x をみて v の次の遷移可能性を返す
(e -> e' -> e) -> -- 緩和の二項演算
e -> -- 初期値 (0, minBound, maxBound, False など)
(v, v) -> -- 状態空間の下界、上界
[(v, e')] -> -- 開始時点の状態
t x -> -- 入力 (時間遷移)
a v e -- Array or UArray
accumArrayDP f op initial (l, u) v0s xs = do
let dp = accumArray op initial (l, u) v0s
foldl' transition dp xs
where
transition dp x =
accumArray op initial (l, u) $
filter (inRange (bounds dp) . fst) $
concatMap (`f` x) (assocs dp)
```
`accumArrayDP`を使うとナップザック問題は以下のような計算の宣言により解くことができる。
```haskell
main :: IO ()
main = do
-- 入力受け取り
[n, w] <- getInts
xs <- replicateM n getTuple
let dp = accumArrayDP @UArray f max minBound (0, wx) [(0, 0)] xs
where
f (w, v) (a, b)
| v == minBound = [] -- 無効なときは状態遷移しない
| otherwise = [(w, v), (w + a, v + b)]
print $ maximum (elems dp)
```
`accumArrayDP` の `where` で定義している関数に着目してほしい。
```haskell
f (w, v) (a, b)
| v == minBound = []
| otherwise = [(w, v), (w + a, v + b)]
```
先に考えた状態遷移関数そのものである。
ナップザック問題に対する DP の計算を、状態遷移のモデルでとらえたことにより「状態遷移関数を渡すだけで問題が解ける」という抽象化に至った。
実装をみてもらえば分かる通り、この実装にはミュータブルな値は一切使っていない。イミュータブルな状態遷移だけでナップザック問題のような DP の問題が解けるというのは、面白いところだと思う。より関数型プログラミングらしい実装になった。
## 競技プログラミングの問題を解く
競技プログラミングにおける DP の問題というのは色々な種類があるように見えるが、計算の構造だけをみると基本的にはそんなに種類はないと思っている。
私のメンタルモデルからするとナップザック問題の亜種か、いわゆる蛙跳び問題の亜種かという区別になる。
蛙跳び問題は、ミュータブルな配列を使わないと実装が難しいが、前者であればこの `accumArrayDP` 関数でおおよそ解くことができている。
いくつか例を挙げる。
### [ABC317 D - President](https://atcoder.jp/contests/abc317/tasks/abc317_d)
比較的最近の DP の問題。これは典型的なナップザック問題である。
```haskell
calcCost :: Int -> Int -> Int
calcCost x y = do
let s = (x + y) `div` 2 + 1
max 0 (s - x)
main :: IO ()
main = do
n <- getInt
xs <- replicateM n $ do
[x, y, z] <- getInts
return (calcCost x y, z)
let sz = sum (map snd xs)
z' = sz `div` 2 + 1
let dp = accumDP @UArray f min maxBound (0, z') [(0, 0)] xs
where
f (w, v) (cost, zi)
| v == maxBound = []
| otherwise = [(min z' (w + zi), v + cost)]
print $ dp ! z'
```
### [ABC266 D - Snuke Panic (1D)](https://atcoder.jp/contests/abc266/tasks/abc266_d)
状態遷移はなにも「選ぶ」「選ばない」の二分岐に限定されるわけではなく、ある状態から三つ以上の複数に遷移しても問題ない。
やはり状態遷移関数を考えることが中心になる。
```haskell
main :: IO ()
main = do
n <- getInt
xs <- replicateM n $ do
[ti, xi, ai] <- getInts
return (ti, xi, ai)
let xs' = zipWith (\(t1, _, _) (t2, x2, a2) -> (t2 - t1, x2, a2)) xs (tail xs)
let dp = accumArrayDP @UArray f max minBound (0, 4) [(0, 0)] $ head xs : xs'
where
f (x, v) (ti, xi, ai)
| v == minBound = []
| otherwise = [(x', v + if x' == xi then ai else 0) | j <- [- ti .. ti], let x' = x + j]
print $ maximum (elems dp)
```
### [ABC312 D - Count Bracket Sequences](https://atcoder.jp/contests/abc312/tasks/abc312_d)
これも比較的最近の問題。括弧の対応を DP で解く。
`(` と `)` のほかに `?` が文字列中に出てきて、`?` は `(`としても `)` としても使って良いという問題設定である。
`?` が出てきたときは、括弧の状態が開き括弧、閉じ括弧両方の状態に遷移するというのが状態遷移関数で宣言的に表現できている。
```haskell
main :: IO ()
main = do
s <- getLine
let n = length s
let dp = accumArrayDP @UArray f addMod (0 :: Int) (0, n) [(0, 1)] s
where
f (_, 0) _ = []
f (w, v) '(' = [(w + 1, v)]
f (w, v) ')' = [(w - 1, v)]
f (w, v) '?' = [(w + 1, v), (w - 1, v)]
f _ _ = error "!?"
print $ dp ! 0
```
### [008 - AtCounter（★4）](https://atcoder.jp/contests/typical90/tasks/typical90_h)
連続しない部分列で A -> T -> C -> O -> D -> E -> R の順に文字が出てくるケースの数え上げ。
A が出てきたら次は T に遷移、T が来たら次は C に遷移、というのを表現できる。
```haskell
data ATCODER = NULL | A | T | C | O | D | E | R deriving (Show, Eq, Ord, Ix)
-- typical90 008 - AtCounter
-- https://atcoder.jp/contests/typical90/tasks/typical90_h
main :: IO ()
main = do
_ <- readLn @Int
s <- getLine
let dp = accumDP @UArray f addMod 0 (NULL, R) [(NULL, 1)] s
where
f (v, x) c = case (v, c) of
(NULL, 'a') -> [(A, x)]
(A, 't') -> [(T, x)]
(T, 'c') -> [(C, x)]
(C, 'o') -> [(O, x)]
(O, 'd') -> [(D, x)]
(D, 'e') -> [(E, x)]
(E, 'r') -> [(R, x)]
_ -> []
print $ dp ! R
```
### [ABC322 E - Product Development](https://atcoder.jp/contests/abc322/tasks/abc322_e)
[Haskell の Array](https://zenn.dev/naoya_ito/articles/87a8a21d52c302) にも書いたが、Haskell の配列の次元は Ix 型で抽象化されていて、多次元配列に表現する場合に実装を変える必要がない。
5次元 DP も `accumArrayDP` で、他の問題同様に実装できる。考えることは状態遷移関数の定義。
```haskell
main :: IO ()
main = do
[n, k, p] <- getInts
xs <- replicateM n getInts
let xs' = map (\x -> x ++ replicate (5 - k) 0) xs
let dp = accumDP @UArray f min (maxBound :: Int) ((0, 0, 0, 0, 0), (p, p, p, p, p)) [((0, 0, 0, 0, 0), 0)] xs'
where
f ((p1, p2, p3, p4, p5), v) [c1, a1, a2, a3, a4, a5]
| v == maxBound = []
| otherwise = [((min p (p1 + a1), min p (p2 + a2), min p (p3 + a3), min p (p4 + a4), min p (p5 + a5)), v + c1)]
f _ _ = error "!?"
let x =
if
| k == 1 -> dp ! (p, 0, 0, 0, 0)
| k == 2 -> dp ! (p, p, 0, 0, 0)
| k == 3 -> dp ! (p, p, p, 0, 0)
| k == 4 -> dp ! (p, p, p, p, 0)
| k == 5 -> dp ! (p, p, p, p, p)
| otherwise -> error "over 5"
print $ if x == maxBound then -1 else x
```
### [D - We Love ABC](https://atcoder.jp/contests/abc104/tasks/abc104_d)
状態遷移が複雑な問題でも、状態遷移にのみ思考を絞り込んで考えることで問題を解くことができるのがこのフレームワークの強みだと思う。
```haskell
data ABC = NULL | A | B | C deriving (Show, Eq, Ord, Ix)
main :: IO ()
main = do
s <- getLine
let dp = accumArrayDP @UArray f addMod 0 (NULL, C) [(NULL, 1)] s
where
f (v, x) u = case (v, u) of
(NULL, 'A') -> [(NULL, x), (A, x)]
(NULL, '?') -> [(NULL, x `mulMod` 3), (A, x)]
(A, 'B') -> [(A, x), (B, x)]
(A, '?') -> [(A, x `mulMod` 3), (B, x)]
(B, 'C') -> [(B, x), (C, x)]
(B, '?') -> [(B, x `mulMod` 3), (C, x)]
(c, '?') -> [(c, x `mulMod` 3)]
(c, _) -> [(c, x)]
let !_ = dbg ("dp", dp)
print $ dp ! C
```
拙作の `accumArrayDP` 関数は、いわゆる「配る DP 」をもとにした状態遷移を実装しているので、もらうDP でないと解くのが難しい問題にはそのままでは適用できない。
具体的には区間の値をまとめてもらってくるような累積和やセグメント木で高速化するDPである。
あとは途中でも述べたとおり、蛙跳びDPや編集距離のように状態空間全体を遷移させるのではない問題にも向いてない。
この辺の見極めは、DP の問題を解いているうちに直感的に判断できるようになったので、特にそれで困っているということはない。
実装的にはいける・いけないが存在しているが、状態遷移モデルで DP を捉えるという考え方はいずれにせよ適用できるので、頭の中ではすべて同じモデルで考えてはいる。
これができるようになってから DP の拡張がずいぶんと楽になり、三次元でも四次元でもビットDPでも、状態遷移モデルに落とし込むことさえできれば無理なく考えられるようになった。
なお DP を関数適用による状態遷移と、二項演算による結合を中心に組み立てたことで代数的な考察もしやすくなっている。このあたりは私もまだ知識不足なので、いつか理解が進んだところで再度まとめてみたいと思っている。