ネタが湧いたので飛び入り参加します。

二次元配列で集めるDP

Haskell Advent Calendar 2024 13日目の記事で nobsun さんが scanArray を提案していた。

scanArray :: (Ix i, Enum i)
          => (a -> b)                 -- 左上の点（直前点なし）
          -> (b -> a -> b)            -- 左辺または上辺上の点（直前点は(i,j-1)または(i-1,j)のどちらか1つ）
          -> (b -> b -> b -> a -> b)  -- 内部の点（直前点は(i-1,j-1)、(i,j-1）、(i-1,j)の3つ
          -> Array (i,i) a
          -> Array (i,i) b
scanArray f g h sa = ta

二次元配列 sa を元の情報として、座標系は右と下に伸びる向きで、sa に関して
上、左、左上の最大3マスからこのマスに向かって集めるDPをする
3方向から集める関数 h は、隣接マスのDP結果と、自分の sa の値から自分のDP結果を求める
集める先が一つ足らない左端と上端、集める先が一つもない左上角について、h の縮退版の g, f も与える

関数を f, g, h と3つも与えるより、Maybe b -> Maybe b -> Maybe b -> a -> b な関数を一つ渡す、だと、ありえない組み合わせが発生して nonexhaustive と怒られるのも嫌で、となると Maybe (Either b (b,b,b)) -> a -> b として、角は Nothing 端は Just (left b) 中は Just (Right (b,b,b)) とでもするか、いっそ専用の3分岐するADTでも定義するか、ということに。それよりは関数3つの方が確かにシンプルなのかも。

既視感

すごく最近どこかで見た気がしたのが、自分の書いた Advent of Code 2024 day 10 の spoiler だった。
当日はもっと書き捨てのコードで答えを出して、後日 spoiler を書くときにかっこつけて、似たようなことをする抽象化を行ったもの。

-- 配列arrの位置iのマスに対して、n i で与えられる周囲のマス j について
-- jとDP結果のタプルのリストjdsから f i jds でDP結果を求める
dp :: Ix i => Array i e -> (i -> [(i, b)] -> b) -> (i -> [i]) -> Array i b
dp arr f n = dpArr
  where
    bnds = bounds arr
    dpArr = listArray bnds
            [ f i [(j, dpArr ! j) | j <- n i]
            | i <- indices arr]

配列の次元数は任意とし、ただし、「隣接するマスとは誰か」も呼び出し元が（関数 n で）与える必要がある。その隣接マスに関するDP結果を貰って集める関数 f は隣接マスのDP値に加えて添え字もリストで受け取る。

なんでこんなことをしたかというと、10日めの問題は二次元配列上の集めるDPではあるが、集める先が左上方向と固定ではなく、4近傍のマスのうち、配列で与えられる標高が「+1高いマス」という、実行時に固定される向きだったため。

配列 arr はnobsunさんの sa と同様に、集める関数に渡すつもりで引数として要求して、元はこんなシグネチャになっていた。

dp :: Ix i => Array i e -> (i -> a -> [(i, a, b)] -> b) -> (i -> a -> [i]) -> Array i b

しかしよく考えてみると arr は呼び出し元も持っているので、必要なら自分でアクセスできる。ということで省いた。n や f の関数がそれを必要とするなら、部分適用で渡せば済む。

さらに今気づいたことに、arr の中身は dp の中では全く触らなくなってしまっている。なのでもっと単純に

-- 添え字範囲bndsの位置iのマスに対して、n i で与えられる周囲のマス j について
-- j と DP 結果のタプルのリスト jds から f i jds でDP結果を求める
dp :: Ix i => (i, i) -> (i -> [(i, b)] -> b) -> (i -> [i]) -> Array i b
dp bnds f n = dpArr
  where
    dpArr = listArray bnds [f i [(j, dpArr ! j) | j <- n i] | i <- range bnds]

これで十分だった。ただ、名前は dp ではひどすぎるのでもうちょっと考え直した方がいい。

使用例

フィボナッチ数列

fib k = dp (0,k) f n ! k
  where
    n i | i < 2     = []
        | otherwise = [i - 1, i - 2]
    f i jds
        | i < 2     = i
        | otherwise = sum $ map snd jds

minPath

minPath sa ij = dp bnds f n ! ij
  where
    bnds = bounds sa
    n (i,j)  = filter (inRange bnds) [(pred i,j), (i, pred j)]
    f ij jds = sa ! ij + minimum (map snd jds)

考察

フィボナッチ数列の例に現れているが、f と n は同じような場合分けをするハメになることがよくある。ひとつの関数に統合してしまうのはどうだろう。

dp :: Ix i => (i, i) -> (i -> ([i], [(i, b)] -> b)) -> Array i b
dp bnds nf = dpArr
  where
    dpArr = listArray bnds
            [ fi [(j, dpArr ! j) | j <- js]
            | i <- range bnds
            , let (js, fi) = nf i]

fib k = dp (0,k) nf ! k
  where
    nf i | i < 2     = ([], const i)
         | otherwise = ([i - 1, i - 2], sum . map snd)

minPath sa ij = dp bnds nf ! ij
  where
    bnds = bounds sa
    nf ij = (n ij, f ij)
    n (i,j)  = filter (inRange bnds) [(pred i,j), (i, pred j)]
    f ij jds = sa ! ij + minimum (map snd jds)

配列を張れない空間でも集めるDP

要素が飛び飛びで、出現しない要素にまで Data.Ix で添え字を振ると広大すぎるような種類の問題がある。
mutableな言語なら、マップを用いたメモ化再帰で、出現する要素のみ考慮することができる。

例：ABC275 D

上の「ひとつの関数に統合した版」の発想を配列から Map に拡張して、事前に列挙できないような要素の定義域でも使える、集めるDP関数を考える。

何らかの Ord T であるような型 T に関する再帰関数 f :: T -> S について、

これを変形した f1 と、（変形の方法は後述）
実際に値を求めたい引数 x :: T
から、メモ化再帰をして memoize f1 x == f x が高速に求める

f1 :: (T -> S) -> T -> ([T], S)

f1 loop x = (ys, a) とする。
ys は、引数 x に関する答えを求めるために再帰呼び出しする必要のあるそのときの引数のリスト。これは loop を使わずに計算する。
a は、引数 x に対する実際の答えを計算する。ただしこの計算の中で f を再帰呼び出ししたいときは代わりに loop を呼ぶこと。

memoize :: Ord t => ((t -> s) -> t -> ([t], s)) -> t -> s
memoize fya x = m M.! x
  where
    m = loop M.empty (S.singleton x)
    loop old new
      | S.null new = old
      | otherwise  = loop old1 new1
      where
        (kvs, jss) = unzip [((k,v),js) | k <- S.elems new, let (js, v) = fya (m M.!) k]
        old1 = M.union old $ M.fromList kvs
        new1 = S.fromList $ concatMap (filter (flip M.notMember old1)) jss

x から始めて、全て尽きるまで fya を呼び出して、必要な y を全て求める。
同時に全てのキーに対する答えがマップ m に集められ、ただしその値は、このマップを参照することで計算するように仕向けている。

ABC275 Dで使ってみたら、GHCが更新されたこともあってか、Int 特化な以前の版より高速になった。

追記：unsafePerformIO

禁断の技を使って、裏でこっそりメモをとる、表向き immutable な関数を作るのも試しておこう。

import Data.IORef
import qualified Data.Map as M
import System.IO.Unsafe

memoize :: Ord d => ((d -> a) -> (d -> a)) -> (d -> a)
memoize mf = f
  where
    memo = unsafePerformIO $ newIORef M.empty
    f x = unsafePerformIO $ do
      m <- readIORef memo
      case M.lookup x m of
        Just a -> return a
        Nothing -> do
          let a = mf f x
          modifyIORef' memo (M.insert x a)
          return a

再帰呼び出しをしたいときは、第１引数に渡された関数を呼び出すことで代用する。
ABC275Dの場合はこんな感じ。

f :: (Int -> Int) -> Int -> Int
f _  0 = 1
f mf x = mf (div x 2) + mf (div x 3)