HaskellでABC431を解く

はじめに

ABC431に参加しました。

A問題

頭と体の重さを比較します。

main = do
  [h,b] <- readInts
  if b >= h then
    print 0
  else
    print (h-b)

A問題提出

B問題

現在どのパーツがくっついているかをIntMapで管理しつつ全体の重さを管理します。パーツの総数が100なので重さを参照する際はリストにインデクスアクセスしている。Nが大きくなるとArrayやVectorにしないとTLEしてしまうのでそちらを使ったほうが安心です。

main = do
  x <- readLn :: IO Int
  n <- readLn :: IO Int
  ws <- readInts
  q <- readLn :: IO Int
  ps <- replicateM q readLn :: IO [Int]
  putStr $ unlines $ map show $ solve IS.empty x ws ps

solve im w ws [] = []
solve im w ws (p:ps)
  | IS.member p im = wSub : solve (IS.delete p im) wSub ws ps
  | otherwise = wAdd : solve (IS.insert p im) wAdd ws ps
  where
    wSub = w - ws !! (p-1)
    wAdd = w + ws !! (p-1)

B問題提出

C問題

頭パーツと体パーツをソートしておいてできるだけ小さい頭と小さい体を使うようにする。
体パーツが頭パーツよりも小さい場合はスルーしてトータルでK体ロボットが作れるかを判定する。

main = do
  [n,m,k] <- readInts
  hs <- sort <$> readInts
  bs <- sort <$> readInts
  let cands = solve hs bs
  if length cands >= k then
    putStrLn "Yes"
  else
    putStrLn "No"

solve [] _ = []
solve _ [] = []
solve hhs@(h:hs) bbs@(b:bs)
  | h <= b = (h,b) : solve hs bs
  | otherwise = solve hhs bs

C問題提出

D問題

最初DPぽいなと思いつつ遷移表がかけなかったので安易とは思いつつ
「2^Nの組み合わせを使って頭の重さが半分を超えた場合に枝かりをする」という方針で取り組んだところやはりTLE。。。
終了間際にDPに戻ってコンテスト終了20分後にDPで解けました。

main = do
  n <- readLn :: IO Int
  whbs <- replicateM n $ do
    [w,h,b] <- readInts
    return (w,h,b)
  let wSum = sum $ map (\(w,h,b) -> w) whbs
  let wSum_2 = wSum `div` 2
  print $ maximum $ solve n wSum_2 whbs

solve :: Int -> Int -> [(Int,Int,Int)] -> [Int] -- A.UArray (Int,Int) Int
solve n wSum_2 whbs = runST $ do
  arr <- AM.newArray ((0,0),(n,wSum_2)) (minBound::Int) :: ST s (STUArray s (Int,Int) Int) -- 型に注意
  AM.writeArray arr (0,0) 0 -- DPテーブルの初期値設定
  forM_ [0..n-1] $ \i -> do
    let (w,h,b) = whbs !! i
    forM_ [0..wSum_2] $ \j -> do
      val0 <- AM.readArray arr (i,j)
      when (val0 /= (minBound::Int)) $ do
        when ( j+w <= wSum_2 ) $ do
          val1 <- AM.readArray arr (i+1,j+w)
          AM.writeArray arr (i+1,j+w) $ max val1 (val0 + h)
        when ( j <= wSum_2 ) $ do
          val1 <- AM.readArray arr (i+1,j)
          AM.writeArray arr (i+1,j) $ max val1 (val0 + b)
  forM [0..wSum_2] $ \j -> do
    val <- AM.readArray arr (n,j)
    return val

D問題提出

おわりに

D問題がDPっぽいなと思ったにもかかわらずTLEなる方針を立ててしまったところが敗因です。もうちょっとDP方向で詰めて考えればよかったです。