ABC344Eやりなおし

全然違うことを調べていて

に通りすがったので、自分はどうやったのかふり返ってみる。

• 位置を無限の精度をもつ何かで実装してMapでやろうとして失敗
• Mutable Vector を使ってあまり美しくないコードで時間に間に合わせた

公式解法

さて、もっと素直に、「連想配列のデータ構造を使った双方向リンクリスト」でやればいいはず、だったらしい。つまり、$A_i$ や $x, y$ の値をキーに、「次」と「前」がどの値かを取り出せる IntMap を維持すればいい。

結果

abc344e :: Int -> [Int] -> Int -> [[Int]] -> [Int]
abc344e _n as _q qus = ans
  where
    as1 = as ++ [0]
    next0 = IM.fromList $ zip (0:as) as1
    prev0 = IM.fromList $ zip as1 (0:as)
    step (next, prev) (1:x:y:_) = (next1, prev1)
      where
        z = next IM.! x
        w = prev IM.! z
        next1 = IM.insert x y $ IM.insert y z next
        prev1 = IM.insert z y $ IM.insert y x prev
    step (next, prev) (2:x:_) = (next1, prev1)
      where
        z = next IM.! x
        w = prev IM.! x
        next1 = IM.insert w z next
        prev1 = IM.insert z w prev
    (nextQ,_) = foldl' step (next0, prev0) qus
    ans = takeWhile (0 /=) $ tail $ iterate (nextQ IM.!) 0

1889msとギリギリではあるけど、普通に間に合った。

以前のコードなどを見てみると、ここで「IntMap でのロスを防ぎたい」とばかりに、

IntMap ( Int  -- 次の値
       , Int) -- 前の値

のようにひとつのマップに詰め込んだために、逆に要素が大きなものになって書き換えが遅くなっているのが敗因な気がする。上のコードは、キーは同じでも

next :: IntMap Int -- 次の値
prev :: IntMap Int -- 前の値

と別のマップにして、値をprimtiveにしておいた方がいい、と。

別の視点から考えると、キーは共通になるけれど、更新が必要になるタイミングが違って、ペアの片方は変更せずに戻す、という操作をするハメになるのは、一緒にするべきでないという警告だったように、今になって思う。

配列を使う解法

前に示した、Mutable Vector を使った間に合う解 も、ひとつの STArray に (Int,Int,Int) （次のアドレス、前のアドレス、このセルのAiの値）の3項組を押し込んで、可読性のかけらもないものだった。3つの配列に分けて書き直してみる。

結果

• Aiがどこにあるかを持つIntMapは、最初のAiの分だけで順に構築し、クエリで更新する。
• その位置にある値が何かの逆Mapは、配列で逐次構築する。
• 順方向リンク(next)、逆方向リンク(prev)、位置から値の逆マップ(val) を STArray でmutableに更新し、Aiからの位置の順マップはimmutableな状態として foldM で伝播させる
• 値0は両端を表す特別な値
• STArray は $1 + N + Q$ 要素を固定で割り当てることで、値が重複しても再利用しないことで富豪的に手間を省いている。具体的には、クエリ1のときにyが既出かどうかを気にせず、自分のクエリ番号から決まるアドレスで決め打ちする。

import qualified Data.IntMap as IM
import Data.Array.ST
import Control.Monad.ST

abc344e :: Int -> [Int] -> Int -> [[Int]] -> [Int]
abc344e n as q qus = runST $
  do
-- 値配列は、0～Nに0,Aiを入れておく
    val <- newArray_ bnds
    forM_ (zip [0 ..] $ 0 : as) (uncurry (writeArray val))
-- 次配列は、0～N-1に(+1)を入れておくあとNは0
    next <- newArray_ bnds
    forM_ (zip [0 ..] [1 .. n]) (uncurry (writeArray next))
    writeArray next n 0
-- 前配列は、1～Nに(-1)を入れておく。あと0はN
    prev <- newArray_ bnds
    forM_ (zip (0 : [1 .. n]) (n : [0 ..])) (uncurry (writeArray prev))
-- クエリを順に処理
-- a2iマップを状態として順に送る mutable な配列は勝手に書き換える
-- 各クエリkは、自分の使うマスがN+kと決められている
    foldM_ (step val next prev) a2i0 $ zip [n + 1 .. ] qus
-- 最後に、0からprevを巡って最終状態を復元する
    post val prev [] =<< readArray prev 0
  where
    bnds = (0, n + q)
    a2i0 = IM.fromList $ (0,0) : zip as [1 ..]
    step :: STUArray s Int Int -> STUArray s Int Int -> STUArray s Int Int -> IM.IntMap Int -> (Int,[Int]) -> ST s (IM.IntMap Int)
    step val next prev a2i (j, 1:x:y:_) =
      do
        let i = a2i IM.! x     -- x@i <---------> ?
        k <- readArray next i  -- x@i <---------> k
        writeArray next i j    -- x@i --> ?@j     k
        writeArray next j k    -- x@i --> ?@j --> k
        writeArray prev k j    -- x@i --> ?@j <-> k
        writeArray prev j i    -- x@i <-> ?@j <-> k
        writeArray val j y     -- x@i <-> y@j <-> k
        return $ IM.insert y j a2i
    step val next prev a2i (_, 2:x:_) =
      do
        let i = a2i IM.! x     -- ? <-> x@i <-> ?
        k <- readArray next i  -- ? <-> x@i <-> k
        h <- readArray prev i  -- h <-> x@i <-> k
        writeArray next h k    -- h ----------> k
        writeArray prev k h    -- h <---------> k
        return a2i -- IM.delete x a2i は冗長
    post _ _ ans 0 = return ans
    post val prev ans i = do
      ai <- readArray val i
      h <- readArray prev i
      post val prev (ai:ans) h

微妙に縦長ではあるが、圧倒的な読みやすさと、前回 MVector を用いて927ms だったものが今回 STArray を用いて 755ms と、タイムの更新まで達成できた。