ABC477 A～F をHaskellで

当日はこれのために前乗りしていました。

A - Traffic Light

問題 ABC477A

シグネチャを決める。

abc477a :: Char -- c
        -> Char -- 答え
abc477a 'B' = 'Y'
abc477a 'Y' = 'R'
abc477a 'R' = 'B'

公式解説、switch - caseでなくifなんだ。

B - Standing Outliers

問題 ABC477B

シグネチャを決める。

abc477b :: Int   -- N
        -> Int   -- D
        -> [Int] -- Xi
        -> [Int] -- 答え（個数は別で数える）

O(N^2)

$N \leq 100$ なので総当たりで調べても大したことはない、

-- まちがい
abc477b _ d xs = [i | (i, x) <- zip [1 ..] xs, all (p x) xs]
  where
    p x y = abs (x - y) >= d

と思ったがこれでは確実に自分が距離0にいるのでそれを除く必要があるが、自分と重なっている誰かの存在も気にしないといけないので結構面倒だ。
「距離D未満の相手がちょうど一人（自分）」と条件を読み替えるのがよいか。

abc477b _ d xs = [i | (i, x) <- zip [1 ..] xs, 1 == length (filter (q x) xs)]
  where
    q x y = abs (x - y) < d

比較側にも添え字を付ける手もある。

abc477b _ d xs = [i | ix@(i,_) <- ixs, all (p ix) ixs]
  where
    ixs = zip [1 ..] xs
    p (i,x) (j,y) = abs (x - y) >= d || i == j

リストを [x(1),...,x(k-1)], x(k), [x(k+1),...,x(N)] の3つに分解したものを全列挙してもいい。

abc477b _ dxs =
    [ i
    | (i, xs1, x:xs2) <- zip3 [1 ..] (inits xs) (tails xs)
    , all (p x) xs1
    , all (p x) xs2 ]
  where
    p x y = abs (x - y) >= d

O(N log N)

$X_i$ の昇順に整列すれば、前後の距離を確認するだけでよい。

import Data.List

abc477b :: Int -> Int -> [Int] -> [Int]
abc477b _ d xs = sort
  [ i
  | ((x0,_), (x1,i) : (x2,_) : _) <- zip xis (tails $ tail xis)
  , x1 - x0 >= d
  , x2 - x1 >= d ]
  where
    xis = sort $ zip (- d : 10^9 + d : xs) [-1 ..]

最小値の下と最大値の上に、必ずD以上離れる番兵を置く。
これを minBound maxBound にするとオーバーフローで痛い目を見る。
番号順に答える必要があるので最後に sort する。

C - Range Search Query

問題 ABC477C

シグネチャを決める。
文字列が長いので ByteString で扱う。

import qualified Data.ByteString.Char8 as BS

abc477c :: Int           -- Q
        -> BS.ByteString -- S
        -> BS.ByteString -- T
        -> [[Int]]       -- query_i
        -> [Bool]        -- 答え

$S$ の $i$ 文字目からが $T$ と一致しているなら1、という個数を数える配列 $C[\cdot]$ があるとする。
クエリ $(L,R)$ は、（正常系な $L + |T| - 1 \leq R$ なとき）
$C[L]$～$C[R - |T| + 1]$ に1が存在するならTrueと読み替えられる。
さらに、$C[\cdot]$ の累積和を $A[\cdot]$ とすると、$A[L-1] < A[R-|T|+1]$ と読み替えられる。

$|T| \leq 10$ なので、$C[\cdot]$ は素朴に計算して $O(|S||T|) \leq 4 \times 10^6$ で問題ない。

常識的なクエリは LとR の間に T が収まる $(L + |T| - 1 \leq R)$ のでこれでいいのだが、そうでない意地悪なクエリも考えられる。例2の2,4にある 5 6 や 6 7 がそれにあたる。
このとき $R - |T| + 1 < L$ と $A[\cdot]$ を読み出す位置関係が逆転するので、偶然条件が満たされてしまう心配はない。

ただし、クエリが端に寄りすぎているせいで、$A[\cdot]$ の添え字範囲を踏み外す危険がある。このことだけは注意する必要がある。

結果

import qualified Data.Vector.Unboxed as UV
import Data.Bool

abc477c :: Int -> BS.ByteString -> BS.ByteString -> [[Int]] -> [Bool]
abc477c q s t = map f
  where
    v = UV.fromList $ scanl (+) (0 :: Int) $ map (bool 0 1 . BS.isPrefixOf t) $ BS.tails s
    len = BS.length t
    f (l:r:_) = len <= r && v UV.! (r - len + 1) > v UV.! (l - 1)

公式解説では、読み出し位置が逆転する場合全てで、読み出すこともせず False を返す判断をするので、踏み外しの危険も同時に回避している。

D - Masking Tape

問題 ABC477D

シグネチャを決める。
クエリ 1 X を Left (Int) クエリ 2 C を Right (Char) で表すことにする。

abc477d :: Int   -- N
        -> Int   -- Q
        -> [Either Int Char] -- query_i
        -> String -- 答え

クエリ2により、ある時刻に塗装が試みられた、という歴史は全てのマスで共有する。
クエリ1により、タイルでマスクされた、もしくはタイルが除かれた、という歴史は個別に管理する。

現在注目しているマスについて、そのマスにタイルを置いた外したの歴史が得られているとする。
最終的にこのマスにタイルが置かれていない場合を考える。
最後にタイルを外した時刻 t1 と、最後に塗った色を塗った時刻 t2 を比べる。
$t1 < t2$ なら、その塗装は有効なので、その色が答え。
$t1 > t2$ なら、その塗装が有効かどうかは判別できない。
歴史をさかのぼり、t1 で外したタイルを置いた時刻 t0 を調べ、時刻 t0 以降のタイル及び塗装の歴史を全て無視し、それ以前だけの歴史について、この判定をやり直す。

最終的にタイルが置かれていた場合は、そのタイルを置いた時刻を調べ、それ以降の歴史を無視するところからループに入ればよい。

結果

歴史を IntMap で扱うとき、「無視」は split で分割することで扱うと直接的だが、lookupLT で、注目している時刻以前だけを調べるようにするだけでも実現できる。

tiles には時刻の逆順で記録されるので都合がよい。
時刻-1にタイルを外し、時刻0にaで塗ってから始まる。

import qualified Data.IntMap as IM
import Data.Array

abc447d :: Int -> Int -> [Either Int Char] -> String
abc447d n _ qus = map solve $ elems tiles
  where
    tiles = accumArray (flip (:)) [-1] (1,n) [(x, i) | (Left x, i) <- zip qus [1 ..]]
    paints = IM.fromDistinctAscList $ (0, 'a') : [(i, c) | (Right c, i) <- zip qus [1 ..]]
    solve ts
      | even $ length ts = go ts
      | otherwise        = go $ maxBound : ts -- Q+1
    go (t0:t:ts)
      | t < tpl   = c
      | otherwise = go ts
      where
        Just (tpl, c) = IM.lookupLT t0 paints

これは ユーザ解説 by shobonvip と同じアプローチだった。

公式解説のやり方

時計を逆回しにして、色を塗るクエリについて、「未着色かつタイルが置かれていない」マスだけに着色することで、着色操作は N 回しか行わないようにする。

状態として
-「未着色かつタイルが置かれていない」マス番号の集合 S
初期値は、順方向にクエリをなぞることで構成する。

• 着色済みのマス番号に対して、塗られた色のマップ M
  初期値は空

を持ち回す。
逆方向にクエリを順に調べ、

• 1 X について、$X \in S$ ならば S から X を除く。そうでなく $X \not\in DOM(M)$ なら $ に X を加える。
• 2 C について、S の要素全てを C に着色して M を更新し、S を空とする。

命令型配列で実装する方が効率的だが、immutableにもできる。

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Functor.Identity

abc447d :: Int -> Int -> [Either Int Char] -> String
abc447d n _ qus = [IM.findWithDefault 'a' x im0 | x <- [1 .. n]]
  where
    sQ = foldl' steps (IS.fromRange (1,n)) qus
    steps is (Right _) = is
    steps is (Left x) = runIdentity $ IS.alterF (Identity . not) x is

    (_,im0) = foldr stepq (sQ, IM.empty) qus
    stepq (Right c) (s, im) = (IS.empty, im1)
      where
        im1 = IM.union im $ IM.fromSet (const c) s
    stepq (Left x) sim@(s, im)
      | IS.member x s = (IS.delete x s, im)
      | IM.notMember x im = (IS.insert x s, im)
      | otherwise = sim

E - Wheel Distance

問題 ABC477E

シグネチャを決める。

abc477e :: Int   -- N
        -> Int   -- Q
        -> [Int] -- Ai
        -> [Int] -- Bi
        -> [[Int]] -- query_i
        -> [Int] -- 答え

$O(N^3)$ なワーシャルフロイド法ではとてもではないが間に合わない。
$O((E+V) \log V)$ なダイクストラ法を Q 回繰り返すのも無理。

1からNのノードがぐるっと輪になり、隣同士に辺が張られている。
中心にノードN+1が陣取り、1からNにスポークのように辺が走っている。

外周を一周したときの距離を $W = \sum A_i$ とする。
中心 N+1 からそれ以外の頂点 v への距離 $d[v]$ をダイクストラ法で求めておく。

クエリ S T が、輪のノードから同じく輪のノードへのものの場合を考える。$(1 \leq S < T \leq N)$
このとき、順方向に輪を辿る経路は $f = \sum_{S \leq k < T} A_i$ 逆方向は $W - f$ である。
中心 N+1 を複数回通るルートは無意味なので、一度だけ通る最短経路は $d[S] + d[T]$ となる。
答えはこれらのうちの最小値である。

$T=N+1$ なクエリの場合は、答えは即座に $d[S]$ である。

結果

$f$ は累積和を用いて $O(1)$ で求める。

import Data.Array.Unboxed
import Data.Array.ST

abc477e :: Int -> Int -> [Int] -> [Int] -> [[Int]] -> [Int]
abc477e n _ as bs = map f
  where
    accs = listArray (0, n) $ scanl (+) 0 as :: UArray Int Int
    b2as = runSTUArray $ dijkstra (1,n+1) g (n+1)
    a = listArray (1,n) as :: UArray Int Int
    b = listArray (1,n) bs :: UArray Int Int
    g k
      | k == n+1  = zip [1 ..] bs
      | otherwise = [(inc k, a ! k), (dec k, a ! dec k), (n+1, b ! k)]
    total = sum as
    inc k
      | k == n    = 1
      | otherwise = succ k
    dec k
      | k == 1    = n
      | otherwise = pred k
    f (s:t:_)
      | t == n + 1 = b2as ! s
      | otherwise  = minimum [dst, dts, ds1t]
      where
        dst = accs ! pred t - accs ! pred s -- S->Tの距離
        dts = total - dst                   -- T->Sの距離
        ds1t = b2as ! s + b2as ! t          -- S->N+1->Tの距離

dijkstra :: (Int, Int)           -- 頂点番号の範囲
         -> (Int -> [(Int,Int)]) -- 隣接頂点とその辺の重み、グラフの情報
         -> Int                  -- 開始点
         -> ST s (STUArray s Int Int)

F - Count Cells in a Window

問題 ABC477F

シグネチャを決める。横着する。

abc477f :: [Int]   -- N,M,Q
        -> [[Int]] -- Li,Ri
        -> [[Int]] -- query_i
        -> [Int]   -- 答え

二次元累積和を使うには $N, M \leq 2 \times 10^5$ は広すぎる。

ダメだったやり方（平方分割）

$N$ 行を $\sqrt N$ 行ごとのブロック $\sqrt N$ 個に分割し、ブロックについてCからDまでの和が累積和の差で求められるようにしておく。
この事前計算は全体で $O(M\sqrt N)$ かかる。

クエリ A,B,C,D について、AからBの範囲のうち、ブロックを完全に覆うものについては事前計算した結果から、ブロックあたり $O(1)$ ブロック最大で $\sqrt N$ 個について求める。
ブロックをかすめる部分については、行ごとに愚直に求める。一行あたり $O(1)$ 行数は最大で $2\sqrt N$
トータルでは $O(\sqrt N)$
クエリ Q 回で $O(Q \sqrt N)$

全体の計算量は $O((M+Q)\sqrt N)$ はかなりギリギリだけど間に合うかと頑張ってみた。
結果 ACx12, TLEx5

公式解説のやり方

配列の区間に対して同じ値を足し込む、また区間の要素の合計を高速に求められるデータ構造、遅延セグメント木を用いる。
行1から行aまでの塗りを足し込んだ遅延セグメント木 $T(a)$ が $0 \leq a \leq N$ まであるとき、
セグメント木Tの要素Cから要素Dまでの合計を $S(T,C,D)$ とすると、
クエリ A,B,C,D は累積和の要領で $S(T(B),C,D) - S(T(A-1),C,D)$ で求められる。

しかし$T(a)$ $(0 \leq a \leq N)$ を全て同時にメモリに置いておくのは難しいので、
$(0 \leq a \leq N)$ を順に構築しつつ、その木を使う必要があるクエリはその木が得られたときに計算をする、
平面走査と呼ぶらしいアイデアでメモリ消費を一行だけにする。

ABC327F で作って放置していた遅延セグメント木の特殊版をもう少しだけ一般化した。
ただしやはり演算は加算、単位元は0、区間更新は同じ値の加算で固定になっている。

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad.ST

abc477f :: [Int] -> [[Int]] -> [[Int]] -> [Int]
abc477f [n,m,q] lrs abcds = runST $
  do
    res <- MUV.replicate q 0 -- クエリは0始まりで扱う
    st <- makeSegTree (m+1)  -- 横方向の位置 C,D,L,R は1始まりで扱う
    forM_ (zip [1 ..] lrs) (\(i, l:r:_) -> do -- 縦方向の位置 A,B は1始まりで扱う
      rangeUpdateSegTree st l (succ r) 1
      forM_ (ab V.! i) (\j -> do
        let jj = if j < 0 then -1-j else j
        let (c,d) = cd UV.! jj
        x <- querySegTree st c (succ d)
        MUV.modify res (if j < 0 then (subtract x) else (x +)) jj
        )
      )
    mapM (MUV.read res) [0 .. q-1]
  where
    cd = UV.fromListN q [(c,d) | _:_:c:d:_ <- abcds]
    ab = V.accum (flip (:)) (V.replicate (n+1) []) $ concat [[(a-1, -1-i), (b, i)] | (i, a:b:_) <- zip [0 ..] abcds]

-- 遅延セグメント木

data SegmentTree s =
 SegmentTree Int       -- 要素数 STreeから復元できるけど
             (STree s) -- キャッシュしている、区間の答え
             (STree s) -- 区間更新での追加の値

type STree s = MUV.MVector s Int

makeSegTree :: Int                  -- 要素数
            -> ST s (SegmentTree s) -- 遅延セグ木

querySegTree
  :: SegmentTree s             -- 区間について問いあわせる
  -> Int  -- 左端（含む）
  -> Int  -- 右端（含まない）
  -> ST s Int -- 結果

-- 区間を最小限で表すwecのノードにxを足し込む
-- 更新されたwecの値をvecの値に重ねて返し、親のvecを更新させる
rangeUpdateSegTree
  :: SegmentTree s
  -> Int  -- 左端（含む）
  -> Int  -- 右端（含まない）
  -> Int  -- 足し込む値
  -> ST s () -- 更新が行われる

提出 1021ms

区間更新の演算が可換であることを仮定しているため、区間更新の情報は wec に一度書き込まれたらそこで確定する。
後から、より狭い区間の更新が重ねられたとき、それは木のさらに奥に浸透してしまう。

可換であることを前提としないようにするには、更新が重ねられたとき（現在のノードの区間更新が id でないとき）、前回の更新内容を両方の子に押し込んでから、到着した新規の更新に対応する形にする必要がある。

G - Frequency Query on Tree

Mo's Algorithm on Trees でググれ、という天啓がTwitter経由で届いたので従った。

• Mo's Algorithm とは、列に対する計算を平方分割して、オフライン処理にして平面走査する感じのアイデア
• 木をdfsで到着順や突入、離脱時刻で列に直すことで、木に対する計算にも適用できる

…大事箱に入れて取っておきます。