HaskellでABC442を解く

はじめに

久しぶりに4問解けましたが解くスピードが遅くて順位はあんまり上がらなかった。

A問題

iかjの文字のみ集計してカウントする。

A問題提出

main = do
  s <- getLine
  let ans = filter (\c -> c =='i' || c == 'j') s
  print $ length ans

B問題

音量とON/OFFの状態を管理する。

B問題提出

main = do
  q <- readLn :: IO Int
  as <- replicateM q readLn
  putStr $ unlines $ map (bool "No" "Yes") $ solve as

solve = step False (0::Int)
  where
    step _playing _vol [] = []
    step playing vol (a:as)
      | a == 1 = f (playing,vol+1) : step playing (vol+1) as
      | a == 2 = let vol1 = if vol > 0 then (vol-1) else 0 in f (playing,vol1) : step playing vol1 as
      | otherwise = let playing1 = not playing in f (playing1,vol) : step playing1 vol as
    f (playing,vol) = playing && vol >= 3

C問題

著者と繋がっている利害関係者の数を管理する。
それ以外の人が査読者としてふさわしいのでその組み合わせを答える。

C問題提出

main = do
  [n,m] <- readInts
  abs <- fmap concat <$> replicateM m $ do
    [a,b] <- readInts
    return [(a,1),(b,1)]
  let g1 = A.accumArray (+) 0 (1,n) abs :: AU.Array Int Int
  putStrLn $ unwords $ map show $ solve g1

solve g = map step [1..n]
  where
    (_,n) = AU.bounds g
    step i = c2 $ n - g A.! i - 1

c2 i
  | i <= 2 = 0
  | otherwise = i * (i-1) * (i-2) `div` 6

D問題

数列の値をゴリゴリ書き換えていく。
また隣がスワップした場合累積和の値が局所的に変更になるので
それも都度書き換えれば良い。
やることは単純だったが配列の値を直接書き換えるようなコードはどうしても長くなってしまい
型エラーでなかなか実装がうまく行かずコード書くのに時間がかかってしまう。

D問題提出

main = do
  [n,q] <- readInts
  as <- readInts
  let arrA = AU.listArray (1,n) as :: AU.UArray Int Int
  let arrS = AU.listArray (0,n) $ scanl (+) 0 as :: AU.UArray Int Int
  qs <- replicateM q readInts
  let cands = solve arrA arrS qs
  let ans = filter (/=(-1)) cands
  putStr $ unlines $ map show ans

solve arrA arrS qs = runST $ do
  arrAA <- AM.thaw arrA :: ST s (STUArray s Int Int)
  arrSS <- AM.thaw arrS :: ST s (STUArray s Int Int)
  forM qs $ \q -> do
    case q of
      (1:x:_) -> do
        -- swap
        val1 <- AM.readArray arrAA x
        val2 <- AM.readArray arrAA (x+1)
        AM.writeArray arrAA x val2
        AM.writeArray arrAA (x+1) val1

        -- scanl
        val3 <- AM.readArray arrSS x
        AM.writeArray arrSS x (val3+val2-val1)
        return (-1)
        
      (2:l:r:_) -> do
        val4 <- AM.readArray arrSS r
        val3 <- AM.readArray arrSS (l-1)
        return (val4-val3)
      _ -> error "error"

おわりに

久しぶりの4問回答はやっぱり嬉しい！