HaskellでABC440を解く

はじめに

今回はC以降すごく難しいと感じました。

A問題

問題文の通り2のy乗で増えていく。

main = do
  [x,y] <- readInts
  print $ x * 2 ^ y

A問題提出

B問題

問題文の通りに並べ替えたりする。

main = do
  n <- readLn :: IO Int
  ts <- readInts
  let ans = map snd $ take 3 $ sort $ zip ts [1..]
  putStrLn $ unwords $ map show $ ans

B問題提出

C問題

なかなか難しかった。解説を読んで

1. 最初に2Wづつで足す
2. 2W*2の幅にする
3. 1つづつずらしながらWの幅で和を取るその中で一番小さいものを解とする
   コンテスト中は1.が思いつかず撃沈しました。解説を読んだあと実装したのが以下です。TLEになってしまいました。

main = do
  t <- readLn :: IO Int
  replicateM_ t $ do
    [_,w] <- readInts
    cs <- readInts
    let ds = map sum $ transpose $ chunksOf (2*w) cs
    let ws = take (2*w) $ ds ++ repeat 0
    let ws2 = ws ++ ws
    print $ minimum $ map (sum . take w) $ take (2*w) $ tails ws2

3.の部分を計算するのに累積和を使うか尺取り方の要領で計算することでO(N)で処理できる。

C問題提出(TLE)

main = do
  t <- readLn :: IO Int
  replicateM_ t $ do
    [_,w] <- readInts
    cs <- readInts
    let ws = A.elems $ AU.accumArray (+) 0 (0,2*w-1) [(i`mod`(2*w),c) | (i,c) <- zip [0..] cs ]
    let ws2 = AU.listArray (0,4*w-1) $ ws ++ ws :: AU.UArray Int Int
    let ps = AU.listArray (0,4*w) $ scanl (+) 0 (AU.elems ws2) :: AU.UArray Int Int
    print $ minimum [ ps AU.! (i+w) - ps AU.! i | i <- [0..2*w-1]]

C問題提出

おわりに

C問題はかなり勉強になった。