ABC289 A～E をHaskellで

ABC288はいつか書きます。

A - flip

問題 ABC289A

文字列で処理してしまおう。

結果

main = getLine >>= putStrLn . map f

f '0' = '1'
f '1' = '0'

B - レ

問題 ABC289B

シグネチャを決める。

abc289b :: Int    -- N
        -> Int    -- M
        -> [Int]  -- ai
        -> [Int]  -- 答え

何だかややこしいことが書いてあるが、$a_i$ と $a_{i+1}$ は一繋がりで、それらを逆順にせよと言っている。

突然だが、foldrでgroupを定義するとき、リストの後ろから積み上げて、リストのリストを作る。

group xs = foldr step [] xs
  where
    step x [] = [[x]]   -- 最後尾の要素
    step x (ys:yss)
      | x == head ys = (x:ys):yss
      | otherwise    = [x]:ys:yss

これと同じ発想で、「レ」で繋がっている区間に分割し、最後に区間ごとに逆転させる。

結果

abc289b n m as = concatMap reverse $ foldr step [[n]] [1 .. pred n]
  where
--    step x [] = [[x]]  -- 空リストは現れない
    step x (ys:yss)
       | elem x as = (x:ys):yss
       | otherwise = [x]:ys:yss

追記：別解

foldl で前から見ていって、「レ」があるならまだ続くから（逆順に）蓄積し、「レ」がないならそこで途切れるから今回の値を追加してから吐き出す、を繰り返すアプローチなら、reverse が不要でこの方が素直な考え方かも。
コードの見た目が少し泥臭いけれど。

abc289b :: Int -> Int -> [Int] -> [Int]
abc289b n m as = loop [] [1..n] as
 
-- loop cs [] _ = cs
loop cs (i:is) (a:as)
  | i == a    = loop (i:cs) is as         -- レがあるなら貯める
  | otherwise = i:cs ++ loop [] is (a:as) -- レがないなら貯めたものを吐き出す
loop cs (i:is) [] = i : cs ++ is

最初に as の末尾に番兵を仕込むと、最後の等式「isは残りがあるのにasは尽きた」を無くせる。
このときはコメントアウトしてある（ある意味で、より自然な）基底部により終了させる。実際にはここの cs は必ず空リストで、asは番兵だけが残った状態になっているはずである。

C - Coverage

問題 ABC289C

シグネチャを決める。$C_i$ は多分使わないので省く。

abc289c :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- aij
        -> Int      -- 答え

何だか大変そうな話に聞こえるが、$1 \leq N \leq 10$ と、考えている集合の規模は大したことがない。また、$1 \leq M \leq 10$ と、考える組み合わせの個数も最大で $2^{10} - 1 = 1023$ とおりと大したことはない。

とはいえこれをそのまま IntSet でするのも芸がないので、競技プログラミングらしく、集合をビット表現で表して、集合の和をビットごとの論理和で計算する。

結果

all1に型注釈をつけることで bs, usの型も定まる。そうでないと不明瞭なままになって怒られる。

import Data.Bits

abc289c :: Int -> Int -> [[Int]] -> Int
abc289c n m ass = length $ filter (all1 ==) us
  where
    bs = map (foldl' setBit 0 . map pred) ass  -- Siのビット表現
    us = foldl' step [0] bs                    -- 「ビット全探索」
    step xs b = map (b .|.) xs ++ xs
    all1 = bit n - 1 :: Int

D - Step Up Robot

問題 ABC289D

シグネチャを決める。
今までのAtCoderの慣例だと、$N,M,X$は1行めでまとめて与えていた気がする。

abc289d :: Int    -- N
        -> [Int]  -- Ai
        -> Int    -- M
        -> [Int]  -- Bj
        -> Bool   -- 答え

手続き的な配列があれば、下の段から順に、0段めからそこに到達できると判明してる、かつそこにモチがないなら、$+A_i$の段へも到達できる、と繰り上げていく $O(NM)$ の手順で計算できる、普通の問題。（いわゆる「配るDP」）

function abc289d(n,as,m,bs,x) {
  const vA = new Array(x+1);  // 到達可能な位置をチェック
  const vB = new Array(x+1);  // モチをマーク
  bs.forEach(b => vB[b] = 1);
  vA[0] = 1;
  for (let k=0; k<=x; k++) {
    if (vB[k] == 1) { continue; }
    if (vA[k] != 1) { continue; }
    as.forEach(a => vA[k + a] = 1; ); // はみ出しチェックは手抜き
  }
  return (vA[x] == 1);
}

Haskellでは、遅延配列を用いて「集めるDP」をする。
つまり、マス$k$は、モチがそこになく、マス$k - A_i$のいずれかが到達可能なら、到達可能である。

結果

import Data.Array

abc289d n as m bs x = vA ! x
  where
    vA = listArray (0,x) $ map f [0..x]  -- 到達可能フラグ
    f 0 = True
    f k | vB ! k = False
    f k = or [vA ! j | j <- takeWhile (0 <=) $ map (k -) as]
    vB = accumArray (||) False (1,x) [(b,True) | b <- bs] -- モチをマーク

E - Swap Places

問題 ABC289E

シグネチャを決める。
問題間に共用できる情報はないので、一つのテストケースについてだけ計算する。

abc289e :: Int          -- N
        -> Int          -- M
        -> [Int]        -- Ci
        -> [(Int,Int)]  -- ui,vi
        -> Int          -- 答え

元のグラフが存在することに意識を引っ張られると、そのグラフ上でのあらゆる経路、とか考えだして訳が分からなくなる。

「高橋君が頂点$A$、青木君が頂点$B$にいる状態」をノードとするグラフを考える。すると

• スタート地点は $(1,N)$
• ゴールは $(N,1)$
• $(A,B)$ から $(C,D)$ に辺がある条件は、
  • 元のグラフで$A$から$C$に移動できる
  • 元のグラフで$B$から$D$に移動できる
  • $C$と$D$は色が異なる

というグラフで、スタート地点からゴール地点までの距離、または到達不可能、を計算すればよいことになる。
辺の重さは一定なので、ダイクストラ法を持ち出さなくても、幅優先探索でするのがよい。Wikipediaにもそう書いてある。

細かい実装の工夫として、グラフの辺を減らすアイデアがある。
最初の状態で二人の足元の色が異なると仮定して(*1)、移動後も異なる色の頂点にいるためには、

• 二人とも、元いた頂点と同じ色の頂点に移動する
• 二人とも、元いた頂点とは異なる色の頂点に移動する

のいずれかでなければならないので、そのような辺だけを持った二つのグラフの重ね合わせで移動を考える。
(*1)について、そもそも頂点$1$と$N$の色が同じ$(C_1=C_N)$だと、どうやってもゴールインできないので、この仮定は妥当である。

高抽象レベル実装

グラフの頂点を (Int,Int) というタプルのままで扱って実装してみる。

import qualified Data.Set as S
import Data.Array

abc289e n m cs uvs
  | cA ! 1 == cA ! n = -1
  | otherwise = loop 0 S.empty (S.singleton (1,n))
  where
    cA = listArray (1,n) cs
    g1 = accumArray (flip (:)) [] (1,n) -- 同色の頂点への辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u == cA ! v, p <- [(u,v), (v,u)]]
    g2 = accumArray (flip (:)) [] (1,n) -- 色の異なる頂点間の辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u /= cA ! v, p <- [(u,v), (v,u)]]
    loop cnt visited news
      | S.member (n,1) news = cnt
      | S.null news         = - 1
      | otherwise           = loop (succ cnt) visited1 news1
      where
        visited1 = S.union visited news
        news1 = S.fromList $
          [ (c,d)
          | (a,b) <- S.elems news
          , c <- g1 ! a, d <- g1 ! b
          , S.notMember (c,d) visited1
          ] ++
          [ (c,d)
          | (a,b) <- S.elems news
          , c <- g2 ! a, d <- g2 ! b
          , S.notMember (c,d) visited1
          ]

ACx53, TLEx11, 129MB という結果に終わった。

ノードに整数を割り当てる

$(a,b)$ に対して $k = a(n+1)+b$ を対応付ける。逆は $(a,b) = \textrm{divMod}(k, n+1)$ で得られる。

import qualified Data.IntSet as IS
import Data.Array

abc289e n m cs uvs
  | cA ! 1 == cA ! n = -1
  | otherwise = loop 0 IS.empty (IS.singleton $ p2i 1 n)
  where
    p2i a b = a * succ n + b
    i2p k = divMod k (succ n)
    cA = listArray (1,n) cs
    g1 = accumArray (flip (:)) [] (1,n) -- 同色の頂点への辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u == cA ! v, p <- [(u,v), (v,u)]]
    g2 = accumArray (flip (:)) [] (1,n) -- 色の異なる頂点間の辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u /= cA ! v, p <- [(u,v), (v,u)]]
    goal = p2i n 1
    loop cnt visited news
      | IS.member goal news = cnt
      | IS.null news        = - 1
      | otherwise           = loop (succ cnt) visited1 news1
      where
        visited1 = IS.union visited news
        news1 = IS.fromList $
          [ cd
          | (a,b) <- map i2p $ IS.elems news
          , c <- g1 ! a, d <- g1 ! b
          , let cd = p2i c d
          , IS.notMember cd visited1
          ] ++
          [ cd
          | (a,b) <- map i2p $ IS.elems news
          , c <- g2 ! a, d <- g2 ! b
          , let cd = p2i c d
          , IS.notMember cd visited1
          ]

ACx63, TLEx1, 26MB と、惜しくも間に合わなかった。

mutable array に手を出す

どうせ整数を割り振っているのだから、visitedをアクセスが$O(\log N)$なIntSetで扱う義理はなく、命令型の書き換え可能配列で実装する。

また、上までは visited は「これ以前に到達した頂点集合」で、news の要素を含まなかったが、今回は news をリストにし、その唯一性を保証することも含めて、visited には news の要素も含まれる形に変わっている。

import Data.Array
import Data.Array.IO

abc289e :: Int -> Int -> [Int] -> [(Int,Int)] -> IO Int
abc289e n m cs uvs
  | cA ! 1 == cA ! n = return $ -1
  | otherwise = do
      visited <- newArray (0, succ n * succ n) False :: IO (IOUArray Int Bool)
      writeArray visited (p2i 1 n) True
      loop 0 visited [(1, n)]
  where
    p2i a b = a * succ n + b
    cA = listArray (1,n) cs
    g1 = accumArray (flip (:)) [] (1,n) -- 同色の頂点への辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u == cA ! v, p <- [(u,v), (v,u)]]
    g2 = accumArray (flip (:)) [] (1,n) -- 色の異なる頂点間の辺のみのグラフ
         [p | (u,v) <- uvs, cA ! u /= cA ! v, p <- [(u,v), (v,u)]]
    goal = p2i n 1
    loop :: Int -> IOUArray Int Bool -> [(Int,Int)] -> IO Int
    loop _ _ [] = return $ -1
    loop cnt visited news = do
      goaled <- readArray visited goal
      if goaled then return cnt else do
        let news11 = [(c,d) | (a,b) <- news, c <- g1 ! a, d <- g1 ! b]
        let news12 = [(c,d) | (a,b) <- news, c <- g2 ! a, d <- g2 ! b]
        news2 <- foldM (\news (c,d) -> do
          let cd = p2i c d
          v <- readArray visited cd
          writeArray visited cd True
          return $ if v then news else (c,d):news
          ) [] (news11 ++ news12)
        loop (succ cnt) visited news2

ACx64, 20MB, 196ms 大勝利。 （※当初、メモリを10MBとtypoしていました。訂正します。）

追記：Data.Vector.Mutableの添え字はIntのみだけど、Data.Array.IOはIxが使えるので、整数への埋め込みをしないように戻した版でもやってみた。
20MB, 204ms と、オーバーヘッドはほとんど無視できる。

コメント

(*1)の前提条件について、グラフを二つに分けるここでのアプローチでは前提条件のチェックが必須で、これをサボるとサンプルの2つめが距離1でゴールしてしまう。
グラフをひとつだけで考え、辿るたびに色が異なることを確認する、普通のアプローチでは、この前提条件を確認しなくても、ゴールに実際入ることができないので問題ない。

本番中、前提条件をチェックしたためにサンプル2がその場ではねられ、「幅優先探索の探索先が空になり、探索がそれ以上進められなくなった場合」のチェックをド忘れして、気づかないままタイムアップしてしまった。痛恨の凡ミス…