ABC403をHaskellで

簡単！簡単！と思ったらコーナーケースの見落としでWAをくらう回だった。
あと謎のTrie推し。
ともあれノーヒントで全ACできためでたい回。（時間内にできる気はしない。）

A - Odd Position Sum

問題 ABC403A

シグネチャを決める。

abc403a :: Int   -- N
        -> [Int] -- Ai
        -> Int   -- 答え
abc403a _n as = sum [a | (a, True) <- zip as $ cycle [True,False]]

B - Four Hidden

問題 ABC403B

シグネチャを決める。

abc403b :: String  -- T
        -> String  -- U
        -> Bool    -- 答え

文字 ? は何とでも等しい文字と見なして、UがTに出現しているかを判定したい。
Char の (==) を newtype と instance で差し替えることで Data.List.isInfixOf が使える。

結果

import Data.List

abc403b :: String -> String -> Bool
abc403b t u = isInfixOf (map Char1 u) (map Char1 t)

newtype Char1 = Char1 Char

instance Eq Char1 where
  Char1 '?' == _ = True
  _ == Char1 '?' = True
  Char1 a == Char1 b = a == b

C - 403 Forbidden

問題 ABC403C

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

abc403c :: [Int]    -- N,M,Q
        -> [[Int]]  -- query_i
        -> [Bool]   -- 答え

ユーザごとに閲覧権限のあるページ番号を二次元のフラグ配列で管理すると、
要素数$NM$が入りきらないし、クエリ2で$M$本のフラグを立てる計算が$Q$回で大変なことになる。

空間の節約のために、ユーザごとの閲覧権限のあるページ番号は IntSet で管理し、
クエリ2が発動しているユーザは特別に印を付けて、IS.fromList [1 .. m] を作らずに済ませる。

結果

IM.IntMap (Maybe IS.IntSet) として、ユーザ番号をキーに、

• キーに要素がないとき、IS.empty が付いていると見なす
• Just IS.IntSet が付いているとき、それは個別の許可の表
• Nothing が付いているとき、クエリ2で権限開放済み

とする方法と、
(IS.IntSet, IM.IntMap IS.IntSet) として、どちらもユーザ番号をキーに、

• 前者に含まれるとき、クエリ2済み
• 後者は個別の許可の表、ないときは空集合を幻視

さらに変種で、後者も IS.IntSet で持ち、Data.Ix で ((1,1),(n,m)) を整数に投影した値で管理することも考えられる。

最後の版で。

import qualified Data.IntSet as IS
import Data.Ix

abc403c :: [Int] -> [[Int]] -> [Bool]
abc403c [n,m,_q] qus0 = loop (IS.empty, IS.empty) qus0
  where
    idx (_:x:y:_) = index ((1,1),(n,m)) (x,y)
    loop _ [] = []
    loop    (mode2, mode1) (q@(1:_):qus) = loop (mode2, IS.insert (idx q) mode1) qus
    loop    (mode2, mode1) ((2:x:_):qus) = loop (IS.insert x mode2, mode1) qus
    loop mm@(mode2, mode1) (q@(3:_):qus) = ans : loop mm qus
      where
        ans = IS.member x mode2 || IS.member (idx q) mode1

D - Forbidden Difference

問題 ABC403D

シグネチャを決める。

abc403d :: Int   -- N
        -> Int   -- D
        -> [Int] -- Ai
        -> Int   -- 答え
abc403d _n d as = ...

ABC401DのようにDPなしで計算できる問題と考え違いしてひっかかった。

考える

普通の $D>0$ の場合を考える。

$A_i$ の並び順は全く関係ない。値として出現するものはいくつでも出現してよいし、
削除すると決めた値は全て消さなければならない。
そこでまず、$A_i$ のそれぞれの値がいくつ出現しているのかの個数を数える。

import Data.Array

    ub = 10^6
-- Aiの値ごとに個数を数える
    cnt = accumArray (+) 0 (0, ub) [(a, 1) | a <- as]

これを D 個飛ばしで D とおりの列について、
つまり [ [cnt ! i | i <- [r, r + d .. ub]] | r <- [0 .. pred d]] についてDPする。

DPの状態は、ひとつ前の枠の値を残すときと残さないときの2とおりについて、最小の合計コスト cosE, cosN
次のステップでは、注目している枠の値の個数 cnti が入力される。

cnti = 0 のとき、この枠には値を残せないので、支払う追加コストはない。
DPの状態遷移は、
cosE' = ∞
cosN' = min cosE cosN
と定義できる。

cnti > 0 のとき、
前に数がある場合はこの数は残せないので、コスト cnti を支払い、値を残さない状態になる。cosN' ≦ cosE + cnti
前に数がない場合は選択ができ、
今回の数を残すとき、コストを払わず、値を残す。cosE' ≦ cosN
今回の数も残さないとき、コストを払い、値を残さない。cosN' ≦ cosN + cnti
まとめると
cosE' = cosN
cosN' = min (cosE + cnti) (cosN + cnti)

さらに cnti = 0 と cnti > 0 の場合を統合して
cosE' = if cnti == 0 then ∞ else cosN
cosN' = cnti + min cosE cosN
とまとめられる。

例外的な場合

$D=0$ のときは、隣同士については考える必要はないが、同じ値は1度しか出現できないので、
重複した値は全て捨てるその個数を数えることになる。

結果

import Data.Array

abc403d :: Int -> Int -> [Int] -> Int
abc403d _n d as
  | d == 0    = ans0
  | otherwise = ans1
  where
    ub = 10^6
-- Aiの値ごとに個数を数える
    cnt = accumArray (+) 0 (0, ub) [(a, 1) | a <- as]
-- D=0のとき、ユニークな要素は一つずつだけ生き残れる
    ans0 = sum [pred c | c <- elems cnt, c > 1]
-- D>0のとき、D飛ばしのDPをD回行う
    ans1 = sum $ map f [0 .. pred d]
    tooBig = div maxBound 2
    f r = uncurry min $ foldl' step (tooBig, 0) [cnt ! i | i <- [r, r + d .. ub]]
    step (cosE, cosN) c =
      ( if c == 0 then tooBig else min cosN (cosE + c)
      , c + min cosN cosE )

$D=0$の場合を見落としていた。
その前に、ABC401Dのように交互に取ればよいだけと勘違いしてWAしていた。

E - Forbidden Prefix

問題 ABC403E

シグネチャを決める。横着して、$T_i$ と $S_i$ を分離しない。

abc403e :: Int      -- Q
        -> [String] -- Ti, Si
        -> [Int]    -- 答え

Trieを使う。

受け付けたYの各文字列yについて、その節点以降にあるyの個数を節点に記録する。
よって毎回の答えはTrieの根にある。

Xを受け付けるとき、その終端位置について、以降にあるyの個数だけ答えが減少する。
また、以降はそこ以降にYの枝を伸ばすことは禁止される。

結果

import Data.Array
import Data.List

data Trie = X                           -- Xで切り取った
          | Trie Int (Array Char Trie)  -- 節点

-- 端を伸ばす処理を入れる代わりに、無限に伸びる空虚な木（immutableならではの技）
leaf :: Trie
leaf = Trie 0 $ listArray ('a','z') $ repeat leaf

-- Xの要素を追加し、そこより下の枝を切る
-- 修正されたYの要素数を根まで反映させる
insertX :: String -> Trie -> Trie
insertX s t =
  case inner t s of
    Nothing     -> t
    Just (t1,_) -> t1
  where
    inner :: Trie
          -> String            -- X
          -> Maybe (Trie, Int) -- 修正した木と、減ったYの個数
    inner X _ = Nothing -- 先行するXで切れた枝を伸ばす必要はない このとき木は変化しない
    inner (Trie sz _)    ""    = Just (X, sz)
    inner (Trie sz arr) (x:xs) =
      case inner (arr ! x) xs of
        Nothing     -> Nothing
        Just (t1,c) -> Just (Trie (sz - c) (arr // [(x, t1)]), c)

-- Yの要素を追加する
-- Yの要素数を根まで反映させる
insertY :: String -> Trie -> Trie
insertY s t = fromMaybe t $ inner t s
  where
    inner :: Trie
          -> String     -- Y
          -> Maybe Trie -- 修正した木
    inner X _ = Nothing  -- 伸ばせなかったときは木を変更しない
    inner (Trie sz arr) "" = Just $ Trie (succ sz) arr
    inner (Trie sz arr) (x:xs) =
      case inner (arr ! x) xs of
        Nothing -> Nothing
        Just t1 -> Just $ Trie (succ sz) $ arr // [(x,t1)]

abc403e :: Int -> [String] -> [Int]
abc403e _q tss = mapAccumL step leaf tss
  where
    step t ('1':_:s) = post $ insertX (drop 2 l) t
    step t ('2':_:s) = post $ insertY (drop 2 l) t
    post t@(Trie sz _) = (t, sz)
    post t = (t, 0)               -- 根がXのとき

F - Shortest One Formula

問題 ABC403F

シグネチャを決める。

abc403f :: Int     -- N
        -> String  -- 答え

こういう話は専門なのだが、混乱してWAしてしまった。

repunit な数 (1,11,111,1111) については、そのまま表示するのが最小で間違いないだろう。
これは number であり、同時に <factor>,<term>,<expr> のどれとも見なせる。

そうでない数については、非終端記号の種類に分けてDPで構築する。
つまり、任意の値 k に対して、

• <expr> としての最小の表現
• <term> としての最小の表現
• <factor> としての最小の表現

の全てを作る。repunitについては、全てが show k で済む。
そうでない数について考える。

+ を使って <expr> を作るには、真に k より小さい値 a の最小の <term> 表現と、
b = k - a の <factor> 表現を使う。これはかならず k-1 個作れる。

* を使って <term> を作るには、k の1とkでない約数 a と b = k / a について、
それぞれ <term> 表現と <factor> 表現を使う。これは k が合成数のときだけ作れる。素数だと作れないことに注意。

<factor> を真に小さい値の表現から作る方法はない。

さて、真に小さい値の表現から作る他に、同じ値の別カテゴリ表現から拝借する方法がある。
<term>は<expr>であり、<factor>は<term>である。
しかし<expr>を<factor>にするには括弧でくくるぶん大きくなる。

これを踏まえると、
<expr> は、上記のように真に小さい値から作ったものか、
（存在するなら）<term> を真に小さい値からつくったものか、いずれかの小さい方である。
<factor>は<expr>をカッコするのでより大きくなるから計算に入れる必要はない。

<term> は、（存在するなら）真に小さい値から作ったものか、
<factor>、これは<expr>をカッコしたもの、のいずれか小さい方である。
このとき、<factor>を作るための<expr>に今考えている<term>が循環しそうだが、
それはカッコのぶん大きくなるので入れる必要はなく、
<expr>のうち、真に小さい値から作ったものだけを考慮すればよい。

<factor> は常に、最終的に選ばれた <expr> の最小値をカッコしたものである。
この区分は一見不要で、<expr> をその場でカッコして作ればよいように思うかもしれないが、
そこをケチると (111) のようにrepunitでバグってしまう。

結果

DP配列の要素のリストは、順に <expr>,<term>,<factor>としての最小の表現を持つ。

import Data.Array

abc403f :: Int -> String
abc403f n = show $ minimum $ expr ! n

data Formula = Formula Int ShowS

instance Show Formula where
  show (Formula _ ss) = ss ""

instance Eq Formula where
  Formula k _ == Formula l _ = k == l

instance Ord Formula where
  compare (Formula k _) (Formula l _) = compare k l

expr :: Array Int [Formula]
expr = listArray bnds $ map f $ range bnds
  where
    bnds = (1, 2000)

    repunits = [1, 11, 111, 1111]

    bra (Formula k ss) = Formula (k + 2) (showChar '(' . ss . showChar ')')

    f n
      | elem n repunits = [rrk, rrk, rrk]
      | null ts         = [e0, brae, brae]
      | otherwise       = [met, min t0 $ bra e0, bra met]
      where
        rrk = Formula (length $ show n) (shows n)
-- Expr
        e0 = minimum
            [ Formula (k + succ l) (s . showChar '+' . t)
            | a <- [1 .. pred n], let [Formula k s,_,_] = expr ! a
            , let b = n - a,      let [_,Formula l t,_] = expr ! b ]
-- Term は空のこともある
        ts = [ Formula (k + succ l) (s . showChar '*' . t)
             | a <- [2 .. pred n]
             , mod n a == 0,    let [_,Formula k s,_] = expr ! a
             , let b = div n a, let [_,_,Formula l t] = expr ! b ]

        brae = bra e0
        t0 = minimum ts
        met = min e0 t0

最初は、それぞれの値について一つだけ最短表現を決めるだけで済む
<factor>として使いたいときはカッコを付ければよいだろう、と早とちりして派手にWAした。

ヒントを見て

フレンズさんいわく

アライグマ「F問題はDPなのだ！　長さ最小のexprとtermを、値が小さい順に求めていけば良いのだ！」

逆に言うと、<factor>は要らないと。
それが欲しい上位の計算で bra $ min expr term で取り出せば済むということかしら。
ただしrepunitのときはカッコをつけてはいけない。関数で包むべきか。

G - Odd Position Sum Query

問題 ABC403G

シグネチャを決める。

abc403g :: Int    -- Q
        -> [Int]  -- yi
        -> [Int]  -- 答え

自分の考えたやり方

セグメント木？

数列Aの並び順は全く関係ない。昇順のBだけ考えればよい。
オフラインで全てのBの要素を先に入手できれば、昇順に背番号を付けて、
それをセグメント木の添字とし、値と演算として
(部分列の要素数, 奇数番目の合計, 偶数番目の合計)
を扱うことができそうだ。
しかし、この問題は一つ前のクエリの結果が次の計算に算入される強制オンライン問題なので、背番号が付けられない。
強引に値の範囲 $[1, 10^9]$ のままでセグメント木を立てることもできない。

動的セグメント木？

値の範囲は $[1,10^9]$ でも要素数は $Q \leq 3\times10^5$ 個しかないので、
二分検索木をバランスさせて作るように、セグメント木をバランスさせて作るか、
バランスを無視してでも、出現した要素に関係する部分だけ枝を実体化させた木として作ることはできないか。

しかし、セグメント木をバランスのために部分的に再計算するのは大変そう。
「出現した部分だけ実体化させる」の方のアイデアを使って、
$[0 < 1,10^9 < 2^{30}-1]$ の範囲の値をそのまま扱う、
具体的には2進表記でビット29から順に、左右に分岐するためにビットの0/1を使う。
これで、30ビット×クエリ数Q以下の節を持つ二分木だけがメモリに作られる。

演算は上で考えたものでよい。

これって

出発点のアイデアはセグメント木だが、これは2進表記に関するTrieに他ならない。
枝刈りもないし分岐は26でなく2だけなので、問題Eよりむしろ単純なのでは。

結果

値の第１要素を「部分列の要素数が奇数である」を表す Bool で単純化した。

import Data.Bits

abc403g :: Int -> [Int] -> [Int]
abc403g _q ys = tail $ map snd $ scanl step (leaf, 0) ys
  where
    step (t, z) y = (t1, ans)
      where
        x = succ $ mod (y + z) (10^9)
        t1 = insertTrie x t
        (_, ans, _) = getVal t1

type T = (Bool,Int,Int)

-- 2進数のTrie
data Trie a = Trie a (Trie a) (Trie a)

getVal :: Trie a -> a
getVal (Trie x _ _) = x

leaf :: Trie T
leaf = Trie (False, 0, 0) leaf leaf

insertTrie :: Int -> Trie T -> Trie T
insertTrie x t = loop 29 t
  where
    loop (-1) (Trie a l r) = Trie (add a (True, x, 0)) l r
    loop i (Trie _a l r)
      | testBit x i = post l (loop (pred i) r)
      | otherwise   = post (loop (pred i) l) r
    post l r = Trie (add (getVal l) (getVal r)) l r

add :: T -> T -> T
add (c1, o1, e1) (c2, o2, e2)
  | c1   = (c1 /= c2, o1 + e2, e1 + o2)
  | True = (c1 /= c2, o1 + o2, e1 + e2)

最初に提出した版, 3973ms, 1005MB は本当にギリギリだった。
上で説明した版, 3071ms, 555MB との相違点は、
Trie の終端を表す Leaf コンストラクタが別にあり、コンストラクタが複数になっていたことと、
タプルの第１要素を要素数のままで扱ったこと。性能向上は多分前者が効いている。

速い解答だと C++ で 108ms, 7MB という次元の違う結果を出しているので、
多分これは出題者的には「こんなことして間に合うとは思ってなかった」想定外解なのでは。

想定解は？

フレンズさんいわく
動的セグ木を使う、貼り付ける値と演算、どちらも自分の考えたものと同じようだ。
普通に想定解だったのか。

公式解説、ユーザ解説によると、さらに効率化するのに平方分割を使うようだ。
それでそんなに変わるものだろうか？