ABC212 A～F をHaskellで

最近のやつが辛いので、8問体制になった最初のところから復習していこうかと思います。

(2023/3/1) Cの計算量が間違っていたので訂正
(2023/3/1) Fの別解を追記

A - Alloy

問題 ABC212A

シグネチャを決める。

abc212a :: Int     -- A
        -> Int     -- B
        -> String  -- 答え

結果

abc212a _ 0 = "Gold"
abc212a 0 _ = "Silver"
abc212a _ _ = "Alloy"

B - Weak Password

問題 ABC212B

シグネチャを決める。

abc212b :: String  -- X1～X4
        -> String  -- 答え

隣同士が全て等しいか、全て「次」が等しいと弱い。

結果

abc212b :: String -> String
abc212b xs
  | cond1 = "Weak"
  | cond2 = "Weak"
  | otherwise = "Strong"
  where
    cond1 = and $ zipWith (==) xs $ tail xs
    cond2 = and $ zipWith (==) (map next xs) $ tail xs

next '9' = '0'
next d   = succ d

C - Min Difference

問題 ABC212C

シグネチャを決める。

abc212c :: Int    -- N
        -> Int    -- M
        -> [Int]  -- Ai
        -> [Int]  -- Bi
        -> Int    -- 答え

Data.IntSet.lookupLT 系の出番。
～LE で以下の最も近い値を探し、～GT でより大きい最も近い値を探す。

結果

Just x <- という生成器で、Nothing な結果がパターンマッチによって捨てられる。

import qualified Data.IntSet as IS

abc212c :: Int -> Int -> [Int] -> [Int] -> Int
abc212c n m as bs = minimum
    [abs (b - x) | b <- bs, Just x <- [IS.lookupLE b s, IS.lookupGT b s]]
  where
    s = IS.fromList as

こちらの計算量は、 lookupの回数が支配的でO(M log N)
N要素のSetを作るのに $O(N \log N)$、これを2M回アクセスするのに $O(M \log N)$、あわせて$O((M+N) \log N)$

高度な機構を使わない解法

IntSet.lookupLE などという強力な武器を持ち出さなくても、$A_i$ と $B_i$ を混ぜてから整列し、隣同士の差を計算すれば答えが求められそう。
ただし、それぞれの値が $A_i$ なのか $B_i$ なのかを忘れてしまうと、$|A_i - A_j|$ など無意味な値が紛れ込んでしまう。そこで、普遍性のある方法としては、区別のためのタグ付けのために、(a_i, False), (b_i, True) のようにする。

import Data.List
import Data.Bits

abc212c :: Int -> Int -> [Int] -> [Int] -> Int
abc212c n m as bs = minimum
  [b - a | ((a,c),(b,d)) <- zip abs (tail abs), c /= d]
  where
    abs = sort $ [(a, False) | a <- as] ++ [(B, True) | b <- bs]

値が整数であることを最大限に活用すると、$A_i * 2$ と $B_i * 2 + 1$ に写すことで、最下位ビットにより見分けるやり方もある。

import Data.List
import Data.Bits

abc212c :: Int -> Int -> [Int] -> [Int] -> Int
abc212c n m as bs = minimum
  [shiftR b 1 - shiftR a 1 | (a,b) <- zip abs (tail abs), testBit a 0 /= testBit b 0]
  where
    abs = sort $ [shiftL a 1 | a <- as] ++ [setBit (shiftL b 1) 0 | b <- bs]

こちらの計算量はソートが支配的で $O((N+M) \log (N+M))$

D - Querying Multiset

問題 ABC212D

状態の変化を伴うクエリに対応する問題なので、シグネチャは決めない。
（表題のとおり）ボールに書きつける数は重複しうるので、IntMapで個数をカウントするなどして多重集合を模倣する必要がある。

しかし、操作2について、IntMap.mapKeys でnaiveに更新するとひどいことになるので、操作2が要求されるたびに、「IntMapのキー $x$ を $x+D$ のことだと見なせ」というずらし量 $D$ を IntMap と共に持ちまわす。次に新たな値 $X_i$ を登録するときは、このずらし量を見越して $X_i - D$ を登録する。

結果

ステップ関数は、毎回のクエリの内容を、現在の状態と共に渡されてそれに対応する。
状態の更新と、必要なら出力を行う。

import qualified Data.IntMap as IM

type State
  = ( IM.IntMap Int  -- ボールに書かれた番号がキー、その番号のボールの個数が値
    , Int)           -- 上のマップのキーに対するオフセット

abc212d :: State     -- 現在の状態
        -> [Int]     -- クエリ
        -> IO State  -- 更新した状態 結果は直接出力
abc212d (m, d) [1,x] = return (IM.insertWith (+) (x-d) 1 m, d)
abc212d (m, d) [2,x] = return (m, d + x)
abc212d (m, d) [3]   = print (x + d) >> return (m1, d)
  where
    (x, c) = IM.findMin m
    m1 = if c == 1 then IM.delete x m else IM.insert x (pred c) m

E - Safety Journey

問題 ABC212E

シグネチャを決める。

abc212e :: Int          -- N
        -> Int          -- M
        -> Int          -- K
        -> [(Int,Int)]  -- Ui, Vi
        -> Int          -- 答え

$N$ 個の都市間の隣接行列を考えると、基本的にオール1、ただし自分への辺がないので対角要素は0、また使えなくなった$M$本の道の要素も0、となる。
これを効率的に$K$乗して、都市1から移動して都市1に戻ってきた場合の数を見れば一瞬で求められるように見えるが、$N \leq 5000$ の $N \times N$ 正方行列を$\log K$ 回掛け算するのはちょっとつらい。

$k$ ステップの遷移後に$1 \sim N$ の都市に到達する場合の数を更新するDPで普通に考える。
ストレートに計算すると、上の隣接行列を一度掛けることと全く同じになる。このとき $N$ 回の乗算と $N-1$ 回の足し算を $N$ 回行うので $O(N^2)$ となる。そしてこれを $K$ ステップ行うので全体では $O(KN^2)$ となる。どこかで節約する必要がある。

道が全て使えるなら、任意の都市へ、あらゆる都市から来られるので、前回の場合の数の総和になる。そこから、自分自身と、使えない道の分をそれぞれ減らすと正しい値になる。この総和を毎回一度だけ計算して、引き算を $N+M$ 回行うと、最初の総和の足し算が $N-1$ 回で、1ステップが $O(N+M)$、全体が $O(K(N+M))$ でできる。

結果

$N$個の足し合わせをしたくらいではオーバーフローしないので、reg を実行するタイミングを切り詰めて高速化している。

import Data.Array.Unboxed

modBase = 998244353
reg x = mod x modBase

type Vec = UArray Int Int

abc212e :: Int -> Int -> Int -> [(Int,Int)] -> Int
abc212e n m k uvs = final ! 1
  where
-- 0ステップ後、都市1にいる場合の数が1、他は0
    initial = listArray (1,n) $ 1 : replicate (pred n) 0
-- 総和から補正で引く必要のある要素リスト
    subs = [(i,i) | i <- [1..n]] ++
           [p | (u,v) <- uvs, p <- [(u,v),(v,u)]]
-- 1ステップの遷移
    step arr = amap reg arr1 :: Vec
      where
        arr1 = accumArray (-) acc (1,n) [(i, arr ! j) | (i,j) <- subs]
        acc = sum $ elems arr
-- 最終状態
    final = iterate step initial !! k

F - Greedy Takahashi

問題 ABC212F

例えば、時刻からそのとき何をしているのかを区間で調べられる IntMap でできた年表を考えて、全てのありうる年表に関して、都市と時刻の区間から、その区間が含まれる年表をひく写像を前処理で作れたなら、それぞれのクエリに関して、都市 $Y_i$ の時刻 $X_i$ を含む年表を取り出し、その時刻 $Z_i$ の状況を答えればできる。後者は $O(\log M)$ くらいになりそうだが、そのような全ての年表と年表への写像を、適度な手間で作る方法が思い当たらない。

どうやら、前処理＋クエリ$Q$回というスタイルではなく、クエリで何を聞かれるのかを意識してバスのデータを処理する、クエリの追加ができないタイプの解き方をするようだ。ということでシグネチャを決める。バスとクエリの情報は手抜きする。

abc212f :: Int      -- N
        -> Int      -- M
        -> Int      -- Q
        -> [[Int]]  -- Ai, Bi, Si, Ti
        -> [[Int]]  -- Xi, Yi, Zi
        -> [[Int]]  -- 答え 都市 [i] または バス [a,b] のリスト

高橋君は、都市$i$ $(1 \leq i \leq N)$ にいるか、バス$j$ $(1 \leq j \leq M)$ に乗っているか、いずれかの場所にいる。この場所を表すID $i$ を、$i \leq N$ なら都市 $i$、$N < i$ ならバス $i - D$ を意味するとする。

バス $(A_i, B_i, S_i, T_i)$ は全部で $M$ 回 $(1 \leq i \leq M)$ 回運行され、時刻 $S_i$ に出発することで、場所 $A_i$ に居た高橋君は場所 $i + D$ に移動する。時刻 $T_i$ に到着することで、場所 $i + D$ にいた高橋君は場所 $B_i$ に移動する。
このバスの運行 $M$ 回で、クエリ $Q$ 個の高橋君の分身 $Q$ 人が、どこに移動していくかを追跡することで、全体で $M$ 回の手間で全てのクエリの答えを求めたい。

そのため、「クエリ番号 $j$ の高橋君 $j$ が、現在どの場所IDに居るか」という写像を、バスの運行に応じて更新する。このとき、高橋君は時刻 $X_j$ に突如として都市 $Y_j$ に出現する。それ以前から都市 $Y_j$ に居ると、バスに乗ってどこかに行ってしまう。また、高橋君ごとに、いつの時刻の居場所を知りたいかは異なる。そこで、バスの発着というイベントに加えて、高橋君 $j$ の出現と、高橋君 $j$ の現在位置を確認するタイミングだというイベントも時系列で管理する必要がある。

また、「高橋君 $j$ が現在どの場所IDに居るか」という写像を、バスの運行によって更新しようとすると、影響を受ける高橋君が誰かを調べるために写像全体を舐めて、その全員の所在を更新することになり、これは $O(Q)$ かかってしまう。
効率化するために、逆写像である「場所ID $p$ に居る高橋君は誰と誰か」という情報を、正方向の写像と一貫性を持たせて同時に更新する。これで舐める必要がなくなる。

しかし、これではまだ不足で、高橋君全員を満載したバスがあちこち走り回ると、結局毎回全員の所在を更新することになる。既に先着の高橋君 $t$ がいる都市に後から到着した高橋君 $s$ はその瞬間、高橋君 $t$ にオンブされ、今後はその所在は高橋君 $t$ の所在に一致する、という、UnionFind的な統合を行うことにする。
オンブされている人は、位置IDに、オンブしてもらっている人の背番号のマイナスを入れて表す。

組み立て

おなじみの import 群。

import qualified Data.IntMap as IM
import Data.Array
import Data.Maybe
import Data.Ord

時刻により起きるイベントは以下の3通り。バスの出発と到着はどちらも Move で表せる。

data Event
  = Move Int Int -- 位置iの人を位置jに移す バスの発着
  | TPop Int Int -- 高橋君iが都市jに出現
  | TEnd Int     -- 高橋君iが任務完了

引数を受け取る。

abc212f :: Int -> Int -> Int -> [[Int]] -> [[Int]] -> [[Int]]
abc212f n m q absts xyzs = answer
  where

イベントの発生タイミング

• 高橋君の出現と消失は、整数の時刻に起きる。
• バスの出発と到着は、+0.5 の時刻に起きて、同時のときには乗り継ぎができる。

を、整数時刻のまま取り扱うのは困難なので、時刻を3倍して、

• 出現と消失は、3の倍数時刻で起きる。
• バスの到着は、3の倍数+1で起きる。
• バスはその直後、3の倍数+2に出発する。

というようにずらして扱う。

    events = map snd $ sortBy (comparing fst) $
-- 高橋君の出現、消失イベント
      [p | (i,(x:y:z:_)) <- zip [1..] xyzs
         , p <- [(x * 3, TPop i y), (z * 3, TEnd i)]] ++
-- バスの発着イベント
      [p | (i,(a:b:s:t:_)) <- zip [n+1..] absts
         , p <- [(s * 3 + 2, Move a i),(t * 3 + 1, Move i b)]]

高橋君は一人もいない、という状態から、イベントを順に消化して状態遷移する。
そのとき同時に、TEndイベントに応じて、結果を出力する。
内容はクエリ番号と、場所IDのタプル。

-- バスの動作をシミュレーション
    (_,irs) = mapAccumL step (IM.empty, IM.empty) events

    type State = (IM.IntMap Int, IM.IntMap Int)
    step :: State -> Event -> (State, Maybe (Int,Int))

移動イベントでは、位置iの高橋君tを位置jに移動させる。位置jに先客sがいればその人に合流する
t2p : Takahashi to Place 負の場所IDは、オンブされている相手のID
p2t : Place to Takahashi 正の高橋IDのみ出現

    step st@(t2p, p2t) (Move i j) =
      case IM.lookup i p2t of
        Nothing -> (st, Nothing) -- 移動する人はいなかった
        Just t ->
          case IM.lookup j p2t of
            Nothing -> ((IM.insert t j t2p, IM.insert j t $ IM.delete i p2t), Nothing) -- 移動先は無人
            Just s  -> ((IM.insert t (-s) t2p, IM.delete i p2t), Nothing) -- 既にsが居る、合流

高橋君iが都市jに出現するイベントでも、先客がいればオンブされる。

    step (t2p, p2t) (TPop i j) =
      case IM.lookup j p2t of
        Nothing -> ((IM.insert i j t2p, IM.insert j i p2t), Nothing) -- 都市jは無人
        Just s  -> ((IM.insert i (-s) t2p, p2t), Nothing) -- 既にsが居る、合流

高橋君iが任務完了。現在位置を報告する。状態は変化しない。
多重にオンブされている可能性があるので、正の位置IDが出てくるまで t2p を手繰る必要がある。

    step st@(t2p, _) (TEnd i)
      | p <= n    = (st, Just (i, [p])) -- 都市ならそのまま
      | otherwise = (st, Just (i, take 2 $ barr ! p)) -- バスならAi,Biに直す
      where
        p = until (0 <) ((t2p IM.!) . negate) (t2p IM.! i)

    barr = listArray (n+1,n+m) absts  -- バスの場所番号 → ABST

最後に、mapAccumL からの出力のうち Just なものだけをID順に整列して完了。

    answer = elems $ array (1,q) $ catMaybes irs

以上のimmutableな解で、2435ms, 238MB でACできた。

同じアイデアを、IntMap の代わりに IOArray で mutable array により実現して、1389ms, 217MB でACした。

感想

どうにか思いついてACできたときはやっぱり楽しい。だけど100分の時間制限に追われながらで、時間内にこれを書けるようになれる気がしない。
バスに乗っていることを表す出力がバス番号でない、一意性の保証されない形なのが、実は全く違う解法のヒントなのだったらどうしよう。

F追記

ふと解説を見てみると、上の心配どおり、これは想定解とは外れていた。想定解に近く、かつHaskellの遅延評価を活かした解答例ができたと思うので追記。

想定解（？）

解説を見ると、gksato氏の解説が、上の、シミュレーション解と同じ考え方。（考え方は同じでも、使っている部品の差なのか、速度がまるで違うの何なん。）

そして、アライさんの解説「1回バスに乗ったら、そのあといつどこにいるか全部決まる」で思いついた別の解法。これなら、クエリが後から来てもいい。

それぞれの都市に対して（つまり都市番号からのArrayで）、どの時刻にどのバスが来るかという、バスの時刻表をIntMapで作る。
クエリの時刻でlookupGEすることで、乗車するバスを見つけることができる。

それぞれのバスに対して（つまりバス番号からのArrayで）、そのバスに乗車したらその後どの時刻にどの場所にいるかを全て記録した年表をIntMapで作る。
バス $(A_i, B_i, S_i, T_i)$ の年表を作るには、後ろから考える。まず降車時刻 $T_i$ 以降に都市 $B_i$ で乗るバスを、上で作った時刻表から探す。そのバスの年表は完成しているので、それの前に、「時刻 $S_i$ にこのバスに乗る」と追記すればよい。時刻 $T_i$ については、「都市 $B_i$ にいる」と書き込みたいが、次のバスがちょうどその時刻に出るならそれは書かないことに注意。

配列を用いた集めるDPと同様に、この二つの時刻表の配列は、遅延評価で、必要なものだけ一度だけ計算されるところがポイント。

この二つの表があれば、クエリ $(X_i, Y_i, Z_i)$ に対して、

1. 都市 $Y_i$ のバス時刻表を見る
2. 時刻 $X_i$ 以降に出発するバスを見つけて、それに乗る
3. そのバスの年表で、時刻 $Z_i$ 以前を見ると、現在位置がわかる

と、IntMapを2度ひくことで答えが得られる。
2の検索に失敗したときは、つまりもうバスは出ないので、都市$Y_i$に留まる。
3の検索に失敗したときは、$Z_i$はバス出発前の時刻なので、やはり都市$Y_i$にとどまっている。
細かい点、3の検索は、+0.5時刻との関係で、lookupLTでする。LEだと誤りとなる。

追記版のコード

-- 前処理で作る、時刻表配列二つ
type Packet =
  ( Array Int (IM.IntMap Int)     -- 都市id -> (時刻 -> バスid)
  , Array Int (IM.IntMap [Int])   -- バスid -> (時刻 -> 位置情報)
  )

-- 前処理
abc212fp :: Int -> Int -> [[Int]] -> Packet
abc212fp n m absts = (cities, buses)
  where
-- Array (都市ID) (IntMap (バスの出発時刻) (バスID))
    cities = fmap IM.fromList $ accumArray (flip (:)) [] (1,n) [(a,(s,i)) | (i,a:b:s:t:_) <- zip [1..] absts]
-- Array (バスID) (IntMap (イベント時刻) (イベント内容)) : 都市iに到着 [i] バスに乗車 [a,b]
    buses = listArray (1,m) $ map getbus absts
    getbus (a:b:s:t:_) = tt2
      where
-- バス到着時刻以降にその都市bから出発するバスがあるなら、そのバスの年表、さもなくば白紙年表
        tt0 = case IM.lookupGE t (cities ! b) of
                Nothing -> IM.empty
                Just (_,j) -> buses ! j
-- 次のバスが時刻tちょうどなら何も書かない、さもなくば、都市[b]に到着と追記
        tt1 = case IM.lookup t tt0 of
                Nothing -> IM.insert t [b] tt0
                _       -> tt0
-- このバスの出発を追記
        tt2 = IM.insert s [a,b] tt1

-- クエリ処理
abc212fm :: Packet -> Int -> Int -> Int -> [Int]
abc212fm (cities, buses) x y z =
  case IM.lookupGE x $ cities ! y of     -- 都市yに時刻x以降に出るバスは
    Nothing -> [y]                       -- ないならyに留まったまま
    Just (_,j) ->                        -- バスjが出るなら乗ってみる
      case IM.lookupLT z $ buses ! j of  -- バスの行動表を見て
        Nothing -> [y]                   -- 何も起きないとは、出発前だった
        Just (_,res) -> res              -- 直近のイベントが答え

結果は737ms, 144MBと、上の解よりかなりよくなった。

lookup 失敗時の処理について、Data.Maybe の関数でもっとコンパクトに書くこともできるが、後で理解できなくなりそうなのでこのぐらいに薄い方が好み。