ABC425をHaskellで

A - Sigma Cubes

問題 ABC425A

シグネチャを決める。

abc425a :: Int -- N
        -> Int -- 答え

定義通りに計算するが、$(-1)^i$ はそうするよりも、negate と id を交互に噛ませる方が計算機に優しそう。

結果

abc425a n = sum [s $ i^3 | (i,s) <- zip [1 .. n] $ cycle [negate, id]]

で、これ何の計算？

$1 \leq N \leq 10$ なら、全部計算して持って置いてもいいけど、それを手で計算したらアホやな。

数列辞典にあった。A232599 一般項も漸化式も書いてある

a(n) = ((-1)^n*(4*n^3+6*n^2-1) +1)/8
a(n) = - 3*a(n-1) - 2*a(n-2) + 2*a(n-3) + 3*a(n-4) + a(n-5)

B - Find Permutation 2

問題 ABC425B

シグネチャを決める。Pが存在しない場合は空リストを返すことにする。

abc425b :: Int -- N
        -> [Int] -- Ai
        -> [Int] -- 答え

1～Nの数のうち、Aに出現しているものはPでもそのまま、Aで-1になっている「空欄」に、残りの数を適当に埋めよ、
ただし、不可能な場合、すなわちAに重複があるときは存在しないと宣言せよ、ということ。

結果

import Data.Array

abc425b :: Int -> [Int] -> [Int]
abc425b n as0
  | any (1 <) $ elems cnt = []                         -- 出現は0か1までしかダメ
  | otherwise             = loop as0 absents
  where
    cnt = accumArray (+) 0 (1,n) [(a, 1) | a <- as0, a > 0] -- Aの各値の出現回数
    absents = [b |(b,0) <- assocs cnt]                      -- 出現0な数リスト
    loop (-1:as) (b:bs) = b : loop as bs           -- -1なところにabsentsを順に当てはめる
    loop ( a:as) bs     = a : loop as bs
    loop []      _      = []

順列組み合わせの総当たり

公式解説では、そういうものを教える教育的配慮なのだろうけど、next_permutation() を使えというオーバースペックなアプローチを紹介している。

import Data.List

abc425b :: Int -> [Int] -> [Int]
abc425b n as = head $ filter match (permutations [1 .. n]) ++ [[]]
  where
    match = and . zipWith prop as
    prop a p = a == -1 || a == p

C - Rotate and Sum Query

問題 ABC425C

シグネチャを決める。

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

位置は0始まりで読み替えることにする。
数列Aの内容自体は変更されないが、クエリ1により、リング状になっていて基準位置がずれると解釈できる。
リングのどの位置からでも連続する区間の和を高速に求めるため、累積和を2周分計算しておく。

結果

import Data.Array.Unboxed

abc425c :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc425c n _q as qus0 = loop 0 qus0
  where
    arr :: UArray Int Int
    arr = listArray (0, n + n) $ scanl (+) 0 $ as ++ as -- 2周分の累積和
    loop :: Int     -- 基準位置
         -> [[Int]] -- クエリ
         -> [Int]   -- 答え
    loop _ase [] = []
    loop base ((1:c:_)  :qus) =       loop base1 qus where base1 = mod (base + c) n
    loop base ((2:l:r:_):qus) = ans : loop base  qus where ans   = arr ! (base + r) - arr ! (base + pred l)

D - Ulam-Warburton Automaton

問題 ABC425D

シグネチャを決める。文字数が多いので ByteString を使う。

import qualified Data.ByteString.Char8 as BS

abc425c :: Int   -- H
        -> Int   -- W
        -> [BS.ByteString] -- Sij
        -> Int -- 答え

指示どおりに命令型の配列を命令型のプログラムで塗っていく。
$10^{100}$回というが、最長で $HW$ ステップあれば真っ黒になってしまう。
といっても、「周囲がちょうど1つ黒な白マスか」を全てのマスについてスキャンすることを繰り返す、という愚直をやると $O((HW)^2)$ になってしまうので、前回塗った黒マスの周囲だけ、を徹底する。

結果

配列の要素を Word8 にしてメモリを節約してみたけれど、多分必要ない。
0～4のとき白マスで、値は周囲の黒マスの個数。
黒マスのとき10。

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Word

abc425d :: Int -> Int -> [BS.ByteString] -> Int
abc425d h w sss = runST $
  do
    fld <- newArray bnds 0 :: ST s (STUArray s (Int,Int) Word8)
    forM_ blacks0 (\ij -> writeArray fld ij 1)
    loop fld blacks0
    length . filter (5 <) <$> getElems fld
  where
    bnds = ((1,1),(h,w))
    blacks0 = [(i, succ j) | (i,ss) <- zip [1 ..] sss, j <- BS.elemIndices '#' ss]
    udlr (i,j) = filter (inRange bnds) [(pred i,j),(succ i,j),(i,pred j),(i,succ j)]

    loop :: STUArray s (Int,Int) Word8 -> [(Int,Int)] -> ST s ()
    loop _ [] = return ()
    loop fld ijs = do
      ijs1 <- filterM (\ij -> do
        c <- readArray fld ij
        when (c == 1) $ writeArray fld ij 10
        return $ c == 1
        ) ijs
      kls <- filterM (\kl -> do
        d <- readArray fld kl
        when (d < 5) $ writeArray fld kl $ succ d
        return $ d == 0) $ concatMap udlr ijs1
      loop fld kls

前のターンで塗ると列挙されたマスに対して、実際に塗ること、塗れたものだけ集めて返すことで重複を排除する、という命令的な副作用で黒塗りをしているのがひとつめの filterM 、
そうして集められた黒く塗られたばかりのマスに対して、その周囲4マスの「周囲の黒の個数」をカウントアップし、次に塗れそうなら集める、というやはり命令的な計算をしているのがふたつめの filterM 我ながらよくこんなHaskellらしくないコードが書けたものだ。

タイトルでググったら、そのまんまのpedia.enがあった。

E - Count Sequences 2

問題 ABC425E

シグネチャを決める。
Mが共通なので、テストケースひとつだけスタイルではなく一括で扱う。

abc425e :: Int -- T
        -> Int -- M
        -> [(Int,[Int])] -- case_i : N, Ci
        -> [Int] -- 各ケースの答え

考える

聞かれていることは、$N$色のボールがそれぞれ$C_i$個ずつあるときに、並べる方法の個数。
高校数学の基本的な問題。
全部で $S = \sum C_i$ 個あるボールのうち、色1を置く位置は ${}_S C_{C_1}$ とおり選ぶことができ、
残った席 $S_1 = S - C_1$ 個から色2を置く位置は ${}_{S_1} C_{C_2}$ とおり選ぶことができ、
…
残った席 $S_{N-1} = C_N$ から色Nを置く位置は ${}_{C_N} C_{C_N} = 1$ 通り、
という掛け算をすればいい。

愚直解、あるいはコンセプトの検証

ただし$S \leq 5000$と少々大きめなので、Integerで愚直に計算すると大変。
理論が正しいか確認するためにやってみる。

abc425e :: Int -> Int -> [(Int,[Int])] -> [Int]
abc425e _t m ncss = [fromIntegral $ f $ map fromIntegral cs | (_n, cs) <- ncss]
  where
    mm = fromIntegral m
    f :: [Integer] -> Integer
    f cs = mod (product $ zipWith comb (scanr1 (+) cs) cs) mm
    comb n k = div (product [succ n - k .. n]) (product [2 .. k])

正しいが、当然間に合わない。

モジュロでの二項係数？

よくある問題だと、合同算術の法が大きな素数に設定されていて、
どの値とも互いに素であることが保証されるため、モジュラ逆数を使って除算も含めて剰余演算でうまいこと comb が計算できる。

しかしこの問題ではMが大きな値であることも素数であることも期待できないので、それもできない。

パスカルの三角形

こういうときは、特に ${}_n C_k$ の n が5000程度なら、パスカルの三角形で事前計算してそれをひけ、とアルゴリズムロジックに書いてある。神託に従う。

import Data.Array.Unboxed
import Control.Monad.ST
import Data.Array.ST

abc425e :: Int -> Int -> [(Int,[Int])] -> [Int]
abc425e _t m ncss = map f ncss
  where
    f (_,cs) = prodd $ zipWith (curry (comb !)) (scanr1 (+) cs) cs

    reg x = mod x m
    add x y = reg $ x + y
    mul x y = reg $ x * y
    prodd :: [Int] -> Int
    prodd = foldl' mul 1

    ub = 5000
    comb :: UArray (Int,Int) Int
    comb = runSTUArray $ do
      arr <- newArray ((0,0),(ub,ub)) 0
      forM_ [0 .. ub] (\i -> do
        writeArray arr (i,0) 1
        writeArray arr (i,i) 1
        forM_ [1 .. pred i] (\j -> do
          x <- readArray arr (pred i, pred j)
          y <- readArray arr (pred i, j)
          writeArray arr (i,j) (add x y)
          )
        )
      return arr

配列の面積が倍というか4倍使う富豪的。
runSTUArray でするこの方法が結局一番速い。いろいろやったけどimmutableなarrayの作り方では追いつけなかった。
公式解説に

簡単に前計算することができます

とあるのに向かって一人で憤慨するなど。

immutableに求めるコードの中で速めでお気に入りなのを一応残しておく。-1の列に0を入れてあるのがミソで、行ごとに計算を分けずに listArray に与える要素を一気に計算している。zwはzipWith'的な。

    comb :: UArray (Int,Int) Int
    comb = listArray ((0,-1),(5000,5000)) combl
    combl = (0 : 1 : replicate 5000 0) ++ zw ad combl (0 : combl)
    ad x y = let z = x + y in if z >= m then z - m else z
    zw f (x:xs) (y:ys) = let z = f x y in z `seq` z : zw f xs ys

さらに速いけどオーバースペックに複雑なアプローチがふたつもユーザ解説で説明されている。

F - Inserting Process

問題 ABC425F

シグネチャを決める。

abc425f :: Int    -- N
        -> String -- T
        -> Int    -- 答え

問題文の「好きな文字」は「任意の文字」と言い換えたい。

文字を追加する方向で考えると、同じ文字列だけど違う位置から取っているものに関して、追加できる先が違ってしまう。
例えば T = "axcyazc" で "ac" とそれより短いものの場合の数を数えると、これは
"axc", "ayc", "azc", "aac", "acc" と色々な相手に配る必要がある。
というめちゃくちゃ遅い解法しか思いつかなくて、フレンズさんに頼る：

アライグマ「F問題はDPなのだ！　
まずは、1文字ずつ消していく操作列を数える問題だと思うのだ。
消し方が違うのに同じ文字列になって困るのは、同じ文字が連続してる箇所でどれを消すかが違うときだけだから、」
「同じ文字が連続してる箇所では一番先頭しか選べない」を条件に入れて、
どの文字が残ってるかでbipDPすればいいのだ！」

消していく方向で考えるなら、T = "axcyaxc" で、同じ "axc" が2箇所から取れても、それらを別に考えてしまえば、文字列として同一なものなのに重複するという無駄がある代わりに、どこから集めればよいかが単純になる。それでいいんだ。

結果

遅延配列による暗黙の集めるDPではきっと遅いので、あきらめて初めからSTを使う。
いわゆる普通の命令型な「ビットDP」になった。

import Data.List
import Data.Array.Unboxed
import Data.Bits

import Data.Array.ST
import Control.Monad
import Control.Monad.ST

abc425f :: Int -> String -> Int
abc425f n t = runST $
  do
    cnt <- newArray_ (0, ub) :: ST s (STUArray s Int Int)
    writeArray cnt 0 1
    forM_ [1 .. ub] (cntf cnt)
    readArray cnt ub
  where
    tA = listArray (0, pred n) t :: UArray Int Char
-- Tの任意の文字を選んだ文字列を、Nビットの整数で表現する

    ub = pred $ bit n

-- xのビットパターンで表されるTに基づく文字列から、
-- いずれか1文字を削除したビットパターンyを作る
-- ただし同じ文字が連続する場合は先頭しか選べない
-- cnt ! y の総和を cnt ! x に書き込んで終わり
    cntf cnt x = writeArray cnt x . summ =<< forM ys (readArray cnt)
      where
        is0 = [i | i <- [0 .. pred n], testBit x i]
        cs = map (tA !) is0
        ys = [clearBit x i | (i,c,d) <- zip3 is0 ('#' : cs) cs, c /= d]

modBase :: Int
modBase =  998244353

reg :: Int -> Int
reg x = mod x modBase
add :: Int -> Int -> Int
add x y = reg $ x + y
summ :: [Int] -> Int
summ = foldl' add 0

$2^{22}$ 要素の配列なんて張れるのかと思ったけど、= 4M で1要素64bit=8byteとして32MB
メモリ制限の1024MiBはまだずっと上か。

多項式時間の解法がユーザ解説で説明されている。

G - Sum of Min of XOR

問題 ABC425G

シグネチャを決める。

abc425g :: Int   -- N
        -> Int   -- M
        -> [Int] -- Ai
        -> Int   -- 答え

考える

Trie

与えられた $x$ に対してxor結果を最小にする$A_i$を選ぶには、なるべく上位の桁が一致するようなものを選ぶ。
単一の $x$ についてそれを探すには、$A_i$ 全員を登録した Trie を作っておき、$x$ とビットが一致するものがある限りそちらへ降りていけばいい。不本意な降り方をした桁のビットが1になった値がxorの最小値となる。

ということでまずはTrieを作る。

data Trie = Leaf | Node Trie Trie

ビット0の桁について分岐した先は、どちらも Leaf にするとその値があるのかないのか見分けがつかなくなるので、ある場合はもう一段（うその）Node を生やすことにする。
Node の第1引数をビット1、第2引数をびっと0とする。なぜかそうしてしまった。

Data.List.partition を使って、注目するビットの0/1で$A_i$を分割していく、割と単純なやり方：

import Data.List
import Data.Bits

-- buildTrie ub as :: ビットub～0で、asのTrieを作る
buildTrie :: Int -> [Int] -> Trie
buildTrie = loop
  where
    loop _ [] = Leaf
    loop (-1) _  = Node Leaf Leaf -- 「あるよ」という印
    loop i xs = Node (loop i1 as) (loop i1 bs)
      where
        i1 = pred i
        (as,bs) = partition (flip testBit i) xs

Trieに対してもう一つ整数を追加する、という操作を繰り返す純粋派なやり方：

buildTrie2 :: Int -> [Int] -> Trie
buildTrie2 ub = foldl trieInsert Leaf . map enBs
  where
    enBs x = [testBit x i | i <- [ub, pred ub .. 0]]
    trieInsert _ [] = Node Leaf Leaf
    trieInsert Leaf bs = trieInsert (Node Leaf Leaf) bs
    trieInsert (Node t1 t0) (True :bs) = Node (trieInsert t1 bs) t0
    trieInsert (Node t1 t0) (False:bs) = Node t1 (trieInsert t0 bs)

二分木である2進数のTrieに対して、多分木で効率的に実装された類似のデータ構造そのものでもあるIntSetを間にはさんで、
Data.List.partition では分割に $O(N)$ かけていたところを $O(\log N)$ でやろうというアプローチ：

import qualified Data.IntSet as IS

buildTrie3 :: Int -> [Int] -> Trie
buildTrie3 ub xs = loop 0 (bit $ succ ub)
  where
    is = IS.fromList xs
    loop p q
      | yes, succ p == q = Node Leaf Leaf
      | yes              = Node (loop m q) (loop p m)
      | otherwise        = Leaf
      where
        m = div (p + q) 2
        yes = Just p <= IS.lookupLT q is -- [p,q) に Ai が含まれる

せっかく作ったので全部示したが、単純な buildTrie が一番速かった。

xの集団がTrieを降りてくる

xを一人一人降ろしていると、$M \leq 10^9$ で間に合わない。集団で扱う必要がある。
集団の人数と開始位置を、2のべきに合わせることで扱いをよくする。

例えば $M$ が2進数で 10100 のとき、
10100 に対して
0xxxx という 0-0000～0-1111 の集団と
100xx という 100-00～100-11 の集団で、[0,10100) は全てになる。
つまり、$M$ の2進数で1なビットについて、

• そのビットより上位桁はそのまま
• そのビットは0に倒す
• そのビットより下位桁は任意、全ての場合を考える

とする。この集団から発生するコストを、エージェントにTrieを辿らせてまとめて求める。

まず前半の確定したビット（上の例の 0- や 100の部分）については、$x$ 全員の意向は一致している。
Trieが一致する方に辿れるならノーコスト、逆に辿るならその桁の重み分、一人分のコストを計算して累積しておく。

後半の xxxx に関しては、0に行きたい $x$ と 1に行きたい $x$ が半数ずつ存在している。
まず後半に入る段階で、$x$ の人数を数えておく。それは現在位置の桁の重みの倍になる。
そして前半の上位桁で確定したコストを人数分支払っておく。

Trieを降りるには、

• 枝が両方に生えているとき、全員が希望する側に進める。再帰的に降りていく計算を二手に分ける、エージェントを分身させることで計算を続行する。
• 枝が片方にしか生えていないとき、希望する側に進めない半数がコストを発生させる。今の桁の重みと、今の人数の半数を掛けた総コストを支払う。
  以後、不本意ながらその半数はもう半数と同行するため、再帰的に木を降りていく計算は枝がある方だけで済むが、人数は半減させない。
• 葉まで到達したらそのエージェントの仕事は完了する。

とする。

import Data.Bool

-- ビット列bs、最上位桁ubなエージェントをtoに放ち、コストを計算する
cost :: [Bool] -> Int -> Trie -> Int
cost bs0 ub to = loop1 0 bs0 ub to
  where
 -- 前半
    loop1 :: Int    -- 累積コスト
          -> [Bool] -- ビット列
          -> Int    -- 桁位置
          -> Trie
          -> Int    -- 答え
    loop1 w [] i t = w * bit (succ i) + loop2 i (bit $ succ i) t
    loop1 w (b:bs) i (Node t1 t0) =
      case tB of
        Node _ _ -> loop1         w    bs (pred i) tB
        _        -> loop1 (setBit w i) bs (pred i) tN
      where
        tB = bool t0 t1 b -- bで行きたい方
        tN = bool t1 t0 b -- でない方
-- 後半
-- Trie全域を覆うxを下らせる。コストを再帰集計する
-- コストは発生した瞬間に外に出す
    loop2 :: Int   -- 桁位置
          -> Int   -- 人数
          -> Trie
          -> Int   -- 答え
    loop2 (-1) _ _ = 0
    loop2 i c (Node t1 t0) =
      case (t1, t0) of
        (Node _ _, Node _ _) -> loop2 (pred i) (div c 2) t1 + loop2 (pred i) (div c 2) t0
        (Node _ _, Leaf)     -> bit i * div c 2 + loop2 (pred i) c t1
        (Leaf, Node _ _)     -> bit i * div c 2 + loop2 (pred i) c t0
        _ -> error "never"

本体

残りの部分。Trieを作り、Mを分解して必要なエージェント全員を出発させる。

abc425g :: Int -> Int -> [Int] -> Int
abc425g _n m as = sum
  [ cost (init bs ++ [False]) ub t
  | bs@(_:_) <- inits mbs, last bs -- Mの前半で、最後が1になっているもの全てについて, 末尾を0にして全域を探索
  ]
  where
    ub = 29
    t = buildTrie ub as
    mbs = [testBit m i | i <- [ub, pred ub .. 0]] -- Mのビット

感想

桁DPでは、上限の値に張り付いている数を追跡し、
注目している桁の数字が1少ない数は、それ以降の桁は0～9任意になるのでその全てをDPで計算する、
というような流れでする。

2進数の桁DPでは、桁の数字が1少ない数としては0しかとれない。

この解法では、

• どのタイミングでエージェントが出発するか
• それらのエージェントが持つ初期コストはいくらか
• 以後の桁について任意になる分について、どうコストを計算するか

を独立に計算したが、いわゆる「桁DP」では、これを全てのエージェントについて重ね合わせて、

• 桁を降りながら必要ならエージェントを出発させる
• 既に出発しているエージェントと上手く合流させて、以降のコスト計算をする

みたいになってるからややこしいのかな、と思った。
普通の桁DPで解ける問題でも、この解法でやったように別々に計算した方が見通しが良くなるかもしれない。
（あるいは逆に、この解法のエージェント全員を統合して通常の桁DPに合わせることで、もっと速くなるかも。）

公式

公式解説では「Trie 木を陽に持たずに再帰的に問題を解く」と、いきなりユーザ解説みたいな話をしている。
Gを解くくらいの人なら Trie はわかるからそれ以外を、ということだろうか。
「陽に持たずに」といって、$B_0, B_1$ で分割しているのは上の buildTrie3 そのものなので、
エージェントが融合して一体となり、さらに Trie の構築と消費が融合したらそうなるのかな、という感じ。