HaskellでEducational DP Contest(A~I)を解く

はじめに

この冬休みを使ってDPコンテストに挑戦しました。
以前はEまでやったことがありましたが今回は頑張ってIまで解きました。
DPの問題をHaskellで解く場合思いつくままにImmutableになんとかしようとするとTLEとかMLEとかになりやすいので結局STモナドを使うのが間違いがないという考えで落ち着いています。

A - Frog1

さてこの問題は配るDPで実装しました。STモナドを使っているので処理するかどうかの場合分けをwhen（やif）で切り分けられます。以下の例はwhenで場合分けしていますがDPテーブルを余裕をもって広く取っておくとwhenでの場合分けをなくすこともできます。

main = do
  n <- readLn :: IO Int
  hs <- readInts
  let arrH = A.listArray (1,n) hs :: A.UArray Int Int
  let ans = solve n arrH
  print ans

solve n arrH = runST $ do
  arrDP <- AM.newArray (1,n) (maxBound::Int) :: ST s (STUArray s Int Int) -- 型に注意
  AM.writeArray arrDP 1 0 -- DPテーブルの初期値設定
  -- 配るDP
  forM_ [1..n-1] $ \i -> do
    cost0 <- AM.readArray arrDP i
    when (i+1 <= n) $ do
      cost1 <- AM.readArray arrDP (i+1)
      let cost = abs $ arrH A.! (i+1) - arrH A.! i
      AM.writeArray arrDP (i+1) $ min cost1 (cost0+cost)
    when (i+2 <= n) $ do
      cost2 <- AM.readArray arrDP (i+2)
      let cost = abs $ arrH A.! (i+2) - arrH A.! i
      AM.writeArray arrDP (i+2) $ min cost2 (cost0+cost)
  AM.readArray arrDP n

提出

B - Frog2

A問題の遷移先2からKに増える部分に対応する。

main = do
  [n,k] <- readInts
  hs <- readInts
  let arrH = A.listArray (1,n) hs :: A.UArray Int Int
  let ans = solve n k arrH
  print ans

solve n k arrH = runST $ do
  arrDP <- AM.newArray (1,n) (maxBound::Int) :: ST s (STUArray s Int Int) -- 型に注意
  AM.writeArray arrDP 1 0 -- DPテーブルの初期値設定
  -- 配るDP
  forM_ [1..n-1] $ \i -> do
    cost0 <- AM.readArray arrDP i
    forM_ [1..k] $ \k -> do
      when (i+k <= n) $ do
        cost1 <- AM.readArray arrDP (i+k)
        let cost = abs $ arrH A.! (i+k) - arrH A.! i
        AM.writeArray arrDP (i+k) $ min cost1 (cost0+cost)
  AM.readArray arrDP n

提出

C - Vacation

遷移先が(A/B/C)の3️⃣つしかないのでImmutableに処理できHaskellで書きやすくてこのタイプのDPは楽しい。

main = do
  n <- readLn :: IO Int
  ((a0,b0,c0):abcs) <- replicateM n $ do
    [a,b,c] <- readInts
    return (a,b,c)
  print $ solve (a0,b0,c0) abcs

solve :: (Int,Int,Int) -> [(Int,Int,Int)] -> Int
solve (a,b,c) [] = maximum [a,b,c]
solve (a0,b0,c0) ((a,b,c):abcs) = solve (a1,b1,c1) abcs
  where
    a1 = max (b0+a) (c0+a)
    b1 = max (a0+b) (c0+b)
    c1 = max (a0+c) (b0+c)

提出

D - Knapsack 1

DPといえばナップサック問題ですね。以下は集めるDPで解いています。
DP配列を２次元にする必要がある。

main = do
  [n,wMax] <- readInts
  wvs <- replicateM n $ do
    [w,v] <- readInts
    return (w,v)
  let arrWV = A.listArray (1,n) wvs :: A.Array Int (Int,Int)
  let ans = solve n wMax arrWV
  print ans

-- 集めるDP
solve n wMax arrWV = runST $ do
  arrDP <- AM.newArray ((0,0),(n,wMax)) (minBound::Int) :: ST s (STUArray s (Int,Int) Int) -- 型に注意
  AM.writeArray arrDP (0,0) 0 -- DPテーブルの初期値設定
  forM_ [1..n] $ \i -> do
    let (w,v) = arrWV A.! i
    forM_ [0..wMax] $ \j -> do
      dp0 <- AM.readArray arrDP (i-1,j)
      if j-w >= 0 then do
        dp1 <- AM.readArray arrDP (i-1,j-w)
        let value1 = max (dp1 + v) dp0
        AM.writeArray arrDP (i,j) value1
      else
        AM.writeArray arrDP (i,j) dp0
  ans <- forM [0..wMax] $ \i -> do AM.readArray arrDP (n,i)
  return $ maximum ans

提出

E - Knapsack 2

ナップサック問題の亜種です。成約からDP配列の横幅をWeightにできないので
取りうるValueの値を横幅にします。

vMax = 100000

main = do
  [n,wMax] <- readInts
  wvs <- replicateM n $ do
    [w,v] <- readInts
    return (w,v)
  let arrWV = A.listArray (1,n) wvs :: A.Array Int (Int,Int)
  let ans = solve n wMax arrWV
  print ans

-- 集めるDP
solve n wMax arrWV = runST $ do
  arrDP <- AM.newArray ((0,0),(n,vMax)) (maxBound::Int) :: ST s (STUArray s (Int,Int) Int) -- 型に注意
  AM.writeArray arrDP (0,0) 0 -- DPテーブルの初期値設定
  forM_ [1..n] $ \i -> do
    let (w,v) = arrWV A.! i
    forM_ [0..vMax] $ \j -> do
      dp0 <- AM.readArray arrDP (i-1,j)
      if j-v >= 0 then do
        dp1 <- AM.readArray arrDP (i-1,j-v)
        if ( dp1 /= (maxBound::Int) && (dp1+w) <= wMax ) then do
          let value1 = min (dp1 + w) dp0
          AM.writeArray arrDP (i,j) value1
        else
          AM.writeArray arrDP (i,j) dp0
      else
        AM.writeArray arrDP (i,j) dp0
  ans <- forM [0..vMax] $ \i -> do
                                    val <- AM.readArray arrDP (n,i)
                                    return (i,val)
  ans2 <- filterM (\(a,b) -> return (b /= (maxBound::Int))) ans
  return $ maximum $ map fst ans2

提出

F - LSC

DPテーブルの処理は手続き型っぽくやる想定であればなんとかなる。

main = do
  s <- getLine
  t <- getLine
  let lenS = length s
  let arrS = AU.listArray (1,lenS) s :: AU.UArray Int Char
  let lenT = length t
  let arrT = AU.listArray (1,lenT) t :: AU.UArray Int Char
  let arrDP = dpSolve arrS arrT
  putStrLn $ reverse $ backTrack arrDP arrS arrT 

backTrack arrDP arrS arrT = step lenS lenT
  where
    (_,lenS) = AU.bounds arrS
    (_,lenT) = AU.bounds arrT
    step 0 _ = []
    step _ 0 = []
    step i j
      | arrS AU.! i == arrT AU.! j = arrT AU.! j : step (i-1) (j-1)
      | arrDP AU.! (i-1,j) >= arrDP AU.! (i,j-1) = step (i-1) j
      | otherwise = step i (j-1)
    

dpSolve arrS arrT = runSTUArray $ do
  let (_,lenS) = AU.bounds arrS
  let (_,lenT) = AU.bounds arrT
  arrDP <- AM.newArray ((0,0),(lenS,lenT)) 0 :: ST s (STUArray s (Int,Int) Int) -- 型に注意
  forM_ [1..lenS] $ \i -> do
    let c1 = arrS AU.! i
    forM_ [1..lenT] $ \j -> do
      let c2 = arrT AU.! j
      if c1 == c2 then do
        val <- AM.readArray arrDP (i-1,j-1)
        AM.writeArray arrDP (i,j) (val + 1)
      else do
        val0 <- AM.readArray arrDP (i-1,j)
        val1 <- AM.readArray arrDP (i,j-1)
        AM.writeArray arrDP (i,j) $ max val0 val1
  return arrDP

提出

G - Longest Path

まずはトポロジカルソート順に点を並べる。次にその順に次に行ける点までの距離をDP配列に保存していく。
DP配列の更新時により長い距離を保存していけば良い。

longestPath :: Int -> A.Array Int [Int] -> [Vertex] -> Int
longestPath n g topo = runST $ do
  dp <- newArray (1,n) 0 :: ST s (STUArray s Int Int)

  forM_ topo $ \u -> do
    du <- readArray dp u
    let vs = g A.! u
    forM_ vs $ \v -> do
      dv <- readArray dp v
      writeArray dp v (max dv (du+1))

  maximum <$> mapM (readArray dp) [1..n]

main = do
  [n,m] <- readInts
  es <- replicateM m $ do
    [a,b] <- readInts
    return (a,b)
  let g = buildG (1,n) es
  let vs = topSort g
  let ans = longestPath n g vs
  print ans

提出

H - Grid 1

経路数を合計していくDP。左上から右下へ横方向あるいは縦方向で合計していく考え方は分かりやすい。

modulo = 1_000_000_007

main = do
  [h,w] <- readInts
  as <- replicateM h getLine
  let grid = AU.listArray ((1,1),(h,w)) $ concat as :: AU.UArray (Int,Int) Char
  let ans = solve grid
  print ans

solve :: AU.UArray (Int,Int) Char -> Int
solve grid = runST $ do
  dp0 <- AM.newArray ((1,1),(h,w)) 0 :: ST s (STUArray s (Int,Int) Int)
  AM.writeArray dp0 (1,1) 1
  forM_ (AU.range ((1,1),(h,w))) $ \(i,j) -> do
    val0 <- AM.readArray dp0 (i,j)
    when (i+1 <= h && grid AU.!(i+1,j) == '.') $ do
      val1 <- AM.readArray dp0 (i+1,j)
      AM.writeArray dp0 (i+1,j) ((val1+val0)`mod`modulo)
    when (j+1 <= w && grid AU.!(i,j+1) == '.') $ do
      val2 <- AM.readArray dp0 (i,j+1)
      AM.writeArray dp0 (i,j+1) ((val2+val0)`mod`modulo)
  AM.readArray dp0 (h,w)

  where
    (_,(h,w)) = AU.bounds grid

  print ans

提出

I - Coins

確率DPは自力ではできなかったので解説を見つつ遷移の仕方を学びました。

main = do
  n <- readLn :: IO Int
  ps <- readDoubles
  let arrP = A.listArray (1,n) ps :: A.UArray Int Double
  let ans = solve n arrP
  let m = (n+1) `div` 2
  let ans2 = [v | i <- [1..m], let v = ans A.! (n+1,i)]
  print $ sum ans2

-- DP[i][j] i枚投げた時に表がj枚となる確率
-- DP[i+1][j] += DP[i][j]*pi
-- DP[i+1][j+1] += DP[i][j]*(1-pi)
solve n arrP = runSTUArray $ do
  arrDP <- AM.newArray ((1,1),(n+1,n+1)) 0 :: ST s (STUArray s (Int,Int) Double) -- 型に注意
  AM.writeArray arrDP (1,1) 1.0
  forM_ [1..n] $ \i -> do
    let pi = arrP A.! i
    forM_ [1..n] $ \j -> do
      dp0 <- AM.readArray arrDP (i,j)
      modifyArray' arrDP (i+1,j) (+dp0*pi)
      modifyArray' arrDP (i+1,j+1) (+dp0*(1-pi))
  return arrDP

modifyArray' :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m ()
modifyArray' ary i f = do
  v <- readArray ary i
  let v' = f v
  v' `seq` writeArray ary i v'

提出

おわりに

J以降は難しすぎて今後やるかどうかはグレーです。多分無理かな。。。