HaskellでABC472を解く

はじめに

今週はD問題まで。C問題で上手いやり方を思いつかず時間がかかってしまったので順位はあまり変わらずでした。

A問題

A以外を.に置き換えます。

A問題提出

main = do
  s <- getLine
  let ans = map (\c -> if c == 'A' then 'A' else '.') s
  putStrLn ans

B問題

左からの累積和と右からの累積和の差の絶対値の中で最大のものが解です。

B問題提出

main = do
  _n <- readLn :: IO Int
  ls <- readInts
  let lefts = scanl (+) 0 ls
  let rights = scanr (+) 0 ls
  let cands = map abs $ zipWith (-) lefts rights
  print $ minimum cands

C問題

リストだけで上手くやるやり方が思いつかず結局Sequenceを使うことにしてACしました。
Sequenceの先頭と末尾を更新しながら条件に合うかどうかを判定します。

C問題提出

main = do
  [_n,m,k] <- readInts
  as <- readInts
  let cands = solve m k as
  forM_ cands $ \i -> do
    putStrLn $ if i then "Yes" else "No"

solve m k = step 0 m empty
  where
    step _sm _cnt _sq [] = []
    step sm cnt sq (a:as)
      | cnt > 0 && sm + a <= k = True : step (sm+a) (cnt-1) (sq |> a) as
      | cnt > 0 = False : step sm (cnt-1) (sq |> 0) as
      | sm-b+a <= k = True : step (sm-b+a) 0 (sq' |> a) as
      | otherwise = False : step (sm-b) 0 (sq' |> 0) as
      where
        b :<| sq' = sq

D問題

爆弾が届かない複数の座乗を始点として多始点BFSで解きます。

D問題提出

type Index = (Int,Int)

main :: IO ()
main = do
  [h,w,k] <- readInts
  grid <- replicateM h getLine
  let arrG = AU.listArray ((1,1),(h,w)) $ concat grid :: AU.UArray Index Char
  -- 始点は縦横に爆弾が無いセル
  let rows = [i | i <- [1..h], all (=='.') $ [arrG AU.! (i,j) | j <- [1..w]]]
  let cols = [j | j <- [1..w], all (=='.') $ [arrG AU.! (i,j) | i <- [1..h]]]
  let starts = [(i,j) | i <- rows, j <- cols]
  print $ length $ filter id $ solve (h,w) k starts arrG

solve :: Index -> Int -> [Index] -> AU.UArray Index Char -> [Bool]
solve (h,w) k starts arrG = runST $ do
  visited <- newArray ((1,1),(h,w)) False :: ST s (STUArray s Index Bool)
  forM_ starts $ \start -> do
    writeArray visited start True -- 最初に訪問済みにする
  bfs visited k starts []
  getElems visited

  where
    bfs _ 0 _ _ = return (-1)
    bfs visited cnt [] [] = return (-1)
    bfs visited cnt [] us = bfs visited (cnt-1) us [] -- 次のグループへ
    bfs visited cnt (v:vs) us = do
      newUs <- foldM push us (dir4 v)
      bfs visited cnt vs newUs
      where
        push acc p
          | not (inRange ((1,1),(h,w)) p) = return acc -- グリッド外の場合
          | arrG AU.! p == '#' = return acc -- 壁の場合
          | otherwise = do
              seen <- readArray visited p
              if seen
                then return acc -- 訪問済みの場合
                else do
                  writeArray visited p True -- enqueue と同時に訪問済み
                  return (p:acc)

dir4 (i,j) = [(i+di,j+dj) | (di,dj) <- [(1,0),(-1,0),(0,1),(0,-1)]]

おわりに

時間かかりすぎたので順位はあまり上がりませんが２週連続D問題まで解けたのは一つの成果です。