HaskellでABC464を解く

はじめに

今週はC問題にハマってしまいました。ハマったと言っても間に合うはずと思っているコードでどうしても1つTLEが取れずでした。コンテスト終了後少しだけ変えて無事ACしました。。。

A問題

'E'と'W’をカウントして多い方を答えます。

A問題提出

main = do
  s <- getLine
  let es = count 'E' s
  let ws = count 'W' s
  putStrLn $ if es > ws then "East" else "West"

count c = length . filter (==c)

B問題

上の行から取れるだけ白の行を取り、その後上下反転して同じことをします。
さらに縦横入れ替えて上の行からと下の行から同じように取り除いたものが解になります。

B問題提出

main = do
  [h,_w] <- readInts
  ccs <- replicateM h getLine
  putStr $ unlines $ solve ccs

solve = step
  where
    step = trimLeftRight . trimTopBottom
    trimTopBottom = reverse . trimTop . reverse . trimTop
    trimLeftRight = transpose . trimTopBottom . transpose
    trimTop [] = []
    trimTop ccs@(c:cs)
      | all (=='.') c = trimTop cs
      | otherwise = ccs

C問題

コンテスト中Mapを使ったやり方でどうしてもTLEの1つが取れず間に合わずでした。
アルゴリズム的にはO(N+M*logM)のはずなんで解けると思ったのですが。。。
コンテスト終了後出力部分をちょっと工夫することでなんとかACできましたがMapをArrayに切り替えるべきでした。

C問題提出

type MultiMap a = MM.Map a Int

fromListMM :: [Int] -> MultiMap Int
fromListMM xs = MM.fromList ixs
  where
    ixs = map (\g -> (head g, length g) ) $ group $ sort xs

insertMM :: Int -> Int -> MultiMap Int -> MultiMap Int
insertMM k c mm = MM.insertWith (+) k c mm -- insertWithを使うとkeyがなかった場合も対応できる

deleteMM :: Int -> Int -> MM.Map Int Int -> MM.Map Int Int
deleteMM k c =
  MM.update f k
  where
    f x =
      let y = x - c
      in if y <= 0
            then Nothing
            else Just y

readInts :: IO [Int]
readInts = unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine

main :: IO ()
main = do
  [n,m] <- readInts
  adbs <- replicateM n $ do
    [a,d,b] <- readInts
    return (a,d,b)
  let as = map (\(a,_,_) -> a ) adbs
  let mm = fromListMM as
  let dabs = sort $ filter (\(_d,(a,b))-> a /= b ) $ map (\(a,d,b) -> (d,(a,b))) adbs
  let dabs' = A.assocs $ A.accumArray (flip (:)) [] (1,m) dabs
  let cands = solve mm dabs'

  BL.putStr $
    BB.toLazyByteString $
      foldMap (\x -> BB.intDec x <> BB.char8 '\n') cands
--  putStr $ unlines $ map show cands

solve mm0 dabs0 = step mm0 dabs0
  where
    step _mm [] = []
    step mm ((_d,abs):dabs) = size : step mm1 dabs
      where
        mm1 = foldl' (\acc (a,b) -> (insertMM b 1 $ deleteMM a 1 acc)) mm abs
        size = MM.size mm1

おわりに

C問題はIntMapで間に合うはずとあれこれ工夫しましたがコンテスト中は間に合わずでした。残念でした。