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方向で詰めて考えればよかったです。

D問題（追記）

最初の方針でも解けたので追記。理論上は最初の方針でもO(NMlogM)なので行けるはずと思ってましたがIntMapではどうしてもダメでした。でArrayだったらもしかして？と思い試したところ無事ACできました。

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -Wname-shadowing #-} -- 内側のスコープの値と同じ名前が外側のスコープにあるとき警告する

import Control.Monad
import qualified Data.ByteString.Char8 as BS
import Data.Char
import Data.List
import qualified Data.Array.Unboxed as AU

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

main :: IO ()
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
  let whbs2 = sortBy (flip compare) whbs
--  print whbs2
  print $ maximum $ AU.elems $ solve wSum_2 whbs2 
--  print $ maximum $ IM.elems $ solve wSum_2 whbs2 

solve :: Int -> [(Int,Int,Int)] -> AU.UArray Int Int
solve wSum_2 = step (AU.listArray (0,wSum_2) (replicate (wSum_2+1) 0) )
  where
    step :: AU.UArray Int Int -> [(Int,Int,Int)] -> AU.UArray Int Int
    step arr [] = arr
    step arr ((w,h,b):whbs) = step arr2 whbs
      where
        cands = [ (xw+yw, x + y) | (xw,x) <- AU.assocs arr, (yw,y) <- [(w,h),(0,b)], xw+yw <= wSum_2]
        arr2 = AU.accumArray max 0 (0,wSum_2) cands

D問題（追記）提出