Haskell: codefes 2017 qualB

A 問題

後ろを drop みたいな関数は見つからなかったので reverse で誤魔化した。大半の haskell の提出もこうやっていた。

main = getLine >>= putStrLn . reverse . drop 8 . reverse

iterate init でいい感じの無限リストを作って 8 番目の要素を取る、という手法も考えられる。

main = getLine >>= putStrLn . (!! 8) . iterate init

一行にこだわらなければ他にも方法はあるだろう。

B 問題

Data.MultiSet は標準にないので atcoder では提出できない。

import Control.Applicative
import Data.MultiSet (MultiSet, isSubsetOf, fromList)

main = do
  _ <- getLine
  a <- fromList . map read . words <$> getLine :: IO (MultiSet Int)
  _ <- getLine
  b <- fromList . map read . words <$> getLine :: IO (MultiSet Int)
  putStrLn . yesno $ b `isSubsetOf` a

yesno :: Bool -> String
yesno True = "YES"
yesno False = "NO"

C 問題

グラフの問題は本当につらい。グラフ自体は Vector に変換しておけばいいんだけど、訪問済み頂点の処理は高価なデータ構造を使わないと純粋関数で書くのは難しい。Map でこの手の操作をすると毎回 time limit ギリギリになる。

import Control.Applicative
import Control.Monad
import Data.List
import Data.Array.IO
import qualified Data.Vector as V
import qualified Data.Map as M
import Debug.Trace
 
main = do
  [n, m] <- map read . words <$> getLine
  g <- newArray (0, n - 1) [] :: IO (IOArray Int [Int])
  forM_ [1..m] $ \_ -> do
    [u, v] <- map pred . map read . words <$> getLine
    readArray g u >>= writeArray g u . (v:)
    readArray g v >>= writeArray g v . (u:)
  g' <- V.fromList <$> getElems g
  let (c, f) = testBipartite g'
  let white = length . filter (==True) . map snd . M.toList $ c
  let black = length . filter (==False) . map snd . M.toList $ c
  print $ if f then white*black - m
               else n*(n-1)`div`2 - m
 
testBipartite :: V.Vector [Int] -> (M.Map Int Bool, Bool)
testBipartite = dfs 0 False M.empty
 
dfs :: Int -> Bool -> M.Map Int Bool -> V.Vector [Int] -> (M.Map Int Bool, Bool)
dfs u c m g = foldl' f (M.insert u c m, True) $ g V.! u
  where
    f :: (M.Map Int Bool, Bool) -> Int -> (M.Map Int Bool, Bool)
    f (m, False) _ = (m, False)
    f (m, True) v = case M.lookup v m
                    of Just c' -> (m, c /= c')
                       Nothing -> dfs v (not c) m g

D 問題

以前は IOArray から IOUArray に書き換えるだけで良かった気がするんだけど、使おうと思うと失敗する。DP なのでおとなしく array を使った。ただ依存関係が狭いので fold でも簡単にいける(書いてないけど)。haskell はメモリ消費量の見積もりが難しい。

import Control.Applicative
import Control.Monad
import qualified Data.Vector as V
import Data.Array.IO

main = do
  n <- readLn :: IO Int
  s <- V.fromList <$> getLine

  dp_1101 <- newArray (0, n) (-inf) :: IO (IOArray Int Int)
  dp_1011 <- newArray (0, n) (-inf) :: IO (IOArray Int Int)
  dp_wait <- newArray (0, n) (-inf) :: IO (IOArray Int Int)

  writeArray dp_wait 0 0

  forM_ [0..n-1] $ \i -> do
    when (i+3 <= n && (V.toList . V.slice i 3) s == "101") $ do
      readArray dp_wait i >>= to dp_1011 (i + 3) 1
      readArray dp_1101 i >>= to dp_wait (i + 3) 1
      readArray dp_wait i >>= to dp_wait (i + 3) 1
    when (s V.! i == '1') $ do
      readArray dp_1011 i >>= to dp_1011 (i + 1) 1
      readArray dp_1101 i >>= to dp_1101 (i + 1) 1
      readArray dp_wait i >>= to dp_1101 (i + 1) 1
      readArray dp_1011 i >>= to dp_wait (i + 1) 1
    readArray dp_wait i >>= to dp_wait (i + 1) 0

  readArray dp_wait n >>= print

to :: IOArray Int Int -> Int -> Int -> Int -> IO ()
to a i c v = readArray a i >>= writeArray a i . max (v + c)

inf :: Int
inf = 10^9