ABC379をHaskellで

A - Cyclic

問題 ABC379A

シグネチャを決める。
$N$は整数として読み込まないべきだろう。

abc379a :: String -- N
        -> String -- 答え
abc379a [a,b,c] = [b,c,a,' ',c,a,b]

B - Strawberries

問題 ABC379B

シグネチャを決める。

abc379b :: Int    -- N
        -> Int    -- K
        -> String -- S
        -> Int    -- 答え

連続する 'O' の個数を $K$ で割った値の合計が答え。

結果

import Data.List

abc379b :: Int -> Int -> String -> Int
abc379b _n k = sum . map (flip div k . length) . filter (('O' ==) . head) . group

C - Sowing Stones

問題 ABC379C

シグネチャを決める。

abc379c :: Int   -- N
        -> Int   -- M
        -> [Int] -- Xi
        -> [Int] -- Ai
        -> Int   -- 答え

$N$ が大きいので、1マスずつちまちま計算する訳にはいかない。

アライさんの「前から丁寧に数える」方法は、場合分けがなんだか面倒くさそう。
フェネックさんの「後ろから調べる」方法を混ぜて、後ろから考えつつ、真面目に数えるやり方を考える。

位置 $p = N + 1$ から始める。$(X_i, A_i)$ を $X_i$ の大きい方から順に考える。

• 位置 $p = 1$ で石がなくなったら、ぴったり成功。
• 位置 $p > 1$ で石がなくなったら、足らずに失敗。
• 最も近い $X_i < p$ に石が $A_i$ 個積まれているとき、これを位置 $p-1, p-2, \dots, p-A_i$ に配って、位置 $p - A_i$ に移動する。
  • ただし、$A_i$ が多すぎて $X_i$ より手前に行ってしまう $p - A_i < X_i$ ときは、石が余って失敗。
  • 大丈夫なとき、$A_i$ 個の石を配るコストは $(p - X_i - 1) + (p - X_i - 2) + \dots + (p - X_i - A_i) = A_i (p - X_i) - A_i (A_i + 1)/2$

結果

import qualified Data.IntMap as IM

abc379c :: Int -> Int -> [Int] -> [Int] -> Int
abc379c n m xs as = loop 0 (succ n) $ IM.toDescList $ IM.fromList $ zip xs as
  where
    loop acc 1 [] = acc -- ぴったり成功
    loop acc _ [] = - 1 -- 足らない
    loop acc p ((x,a):xas)
      | p1 < x    = - 1 -- 多すぎ
      | otherwise = loop (acc + cost) p1 xas
      where
        p1 = p - a
        cost = a * (p - x) - div (a * succ a) 2

D - Home Garden

問題 ABC379D

シグネチャを決める。

abc379d :: Int     -- Q
        -> [[Int]] -- query_i
        -> [Int]   -- 答え

自分はうっかり、植えた日をキー、植わっている個数を値にした IntMap でごりごり計算したけれど、確かに、成長は一定なので、キューで十分だということで、Data.Sequence で書き直し。

結果

import qualified Data.Sequence as Q

abc379d :: Int -> [[Int]] -> [Int]
abc379d _q qis = loop 0 Q.empty qis
  where
    loop _ _ [] = []
    loop d q ([1]  :qis) = loop d (q Q.:|> d) qis
    loop d q ([2,t]:qis) = loop (d + t) q qis
    loop d q ([3,h]:qis) = Q.length q1 : loop d q2 qis
      where
        (q1, q2) = Q.spanl (d - h >=) q

E - Sum of All Substrings

問題 ABC379E

シグネチャを決める。
対象の数 $S$ の桁数 $N$ がとんでもない大きさなのに、求める値はモジュロ演算でなく真値だという。多倍長整数を使って力任せにやる、という問題を出すとは思えないので、桁ごとに何かする方針を想定して、文字列で扱うことにする。

import Data.Char

abc379e :: Int     -- N
        -> String  -- S
        -> String  -- 答え

$S = d_1 d_2 \dots d_N$ として、ある桁 $k$ の数字 $d_k$ が最終結果のどこにどうなるかというと、

(1)
より右にある数字 $k+1, \dots, N$ の個数のそれぞれの位に、何度か足し込まれる。
例えば $S = 12345$ のとき、数字 $3$ は $3,30,300$ として結果に寄与する。

(2)
それぞれ何度寄与するかというと、より左にある数字 $1, \dots, k-1$ の個数のそれぞれの形で、つまり $k$ 回ずつとなる
例えば $S = 12345$ のとき、数字 $3$ は $3,23,123$ として1の位に、$34,234,1234$ として10の位に、$345,2345,12345$ として 100 の位に足し込まれる。

結果をとりあえず、それぞれの位に対する合計を保持する配列 $sum[1 .. N]$ で保持すると考える。
すると、上の(1)(2)より、$d_k$ は それぞれの桁に $k$ 回足し込むので、$sum[k .. N]$ に $k \cdot d_k$ を足し込む。

この配列のそれぞれの要素は、途中で繰り上がりを細かく気にしなくても溢れないので、繰り上がり操作は全てを足し合わせた後ですればよい。

それより、「$sum[k .. N]$ に $k \cdot d_k$ を足し込む。」を素朴にやると $O(N^2)$ の手間が掛かってしまう問題がある。
これは、足し込む区間が必ず $k$ から $N$ なので、累積和をすればよい。

結果

abc379e _n s
  | c /= 0 = shows c $ map intToDigit ds
  | otherwise =        map intToDigit ds -- leading 0 がWAになる…
  where
    (c, ds) = foldr step (0 , []) $ scanl1 (+) $ zipWith (*) [1 ..] $ map digitToInt s
    step x (c, ds) = (q, r : ds)
      where
        (q,r) = divMod (x + c) 10

F - Buildings 2

問題 ABC379F

シグネチャを決める。

abc379f :: Int     -- N
        -> Int     -- Q
        -> [Int]   -- Hi
        -> [[Int]] -- li,ri
        -> [Int]   -- 答え

自分の考えた、割とダイレクトな方法

前処理(1)
右の方から順に考える。
新たにビルが出現すると、それより低いビルは遮られて見えなくなる。
今注目しているビルのすぐ左から見たときに見えるビルの様子を、高さと軒数の IntMap で管理する。
そしてこの IntMap を全ての位置について求めて、配列に入れておく。

前処理(2)
ビルの高さを、max を演算とするセグメント木に入れておく。
後から変更はしないので、immutable arrayで実装できる。

ここまでが事前準備。

個々のクエリ $(l_i,r_i)$ について、$r_i$ から見えるビルとは、
前処理(1)でビル $r_i + 1$ を追加した時点の IntMap im である。これを配列から取り出す。
これらの中でビル$l_i$からも見えるものとは、$l_i + 1$ から $r_i$ の間のビルの高さの最大値以上のビルだけ。
この高さ h は、前処理(2)のセグメント木から取り出せる。
im の h 以上のキーだけを split で取り出し、値を足し合わせたら答え。

import qualified Data.IntMap as IM
import Data.Array

abc379f :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc379f n _q hs lrs = map f lrs
  where
    cntMapArr = listArray (1, succ n) $ scanr step (IM.singleton maxBound 0) hs
    step h im = IM.insert h (succ cnt) $ snd $ IM.split (pred h) im
      where
        Just (_,cnt) = IM.lookupGE h im
    rqa = makeRQArray max 0 hs
    f (l:r:_) = ans
      where
        h = queryRQArray rqa l r
        Just (_, ans) = IM.lookupGE h $ cntMapArr ! succ r

data RQArray a = RQA (a->a->a) a Int (Array Int a)

makeRQArray :: (a->a->a) -> a -> [a] -> RQArray a
makeRQArray op e xs = RQA op e w arr
  where
    len = length xs
    w = until (len <=) (2 *) 1
    arr = listArray (0, w + pred w) $ map f [0 .. w - 2] ++ xs
    f k = op (arr ! (k + succ k)) (arr ! (2 * succ k))

queryRQArray :: RQArray a -> Int -> Int -> a
queryRQArray (RQA op u w arr) a b = loop 0 w 0
  where
    loop p w i
      | q <= a || b <= p = u
      | a <= p && q <= b = arr ! i
      | otherwise        = op l r
      where
        l = loop  p       w2 (i + succ i)
        r = loop (p + w2) w2 (2 * succ i)
        q = p + w
        w2 = div w 2

アライさんの方式

自分のと似ているように見えて、大きく違う点がある。
「見えているビル」は、位置の順に考えたとき、その高さは広義単調増加する。
自分はそれを、見える高さまで切り出すために IntMap を使った。

アライさんの方式では、高さの情報は捨てて、見えているビルをその位置の昇順リストで表現しておく。
$(l_i, r_i)$ というクエリの答えは、$l_i + 1$ のビルを処理した時点での見えるビルリストについて、$r_i$ より大きな位置の要素の個数を数えたらそれが答え、ということ。
配列の二分探索をして、添え字を見たら個数がわかる、というよくある解法だ。

これを再現するために、個数を持った IntMap を個別に作るのは、$O(N^2)$ のカウントアップを引き起こして失敗。
定数は大きくなるが、Int のつかない Map や Set は $O(\log N)$ で背番号がらみの操作が行えるので、こちらを使っていく。

import qualified Data.Set as S
import Data.Array.Unboxed

abc379f :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc379f n _q hs = map f
  where
    hA = listArray (1,n) hs :: UArray Int Int
    cA = listArray (0,n) $ scanr step S.empty [1 .. n] :: Array Int (S.Set Int)
    step i s = S.insert i $ until cond S.deleteMin s
      where
        hi = hA ! i
        cond s = S.null s || hA ! S.findMin s >= hi
    f (l:r:_) =
      case S.lookupGT r cAl of
        Nothing -> 0
        Just j  -> S.size cAl - S.findIndex j cAl
      where
        cAl = cA ! l

あとがき

公式解説だと、cA を配列にとる代わりに、クエリの方を位置ごとにまとめておいて、その位置のマップが決まったら一括で処理するオフラインな計算をする、オンラインでするにはまた違うやり方が必要、みたいに書いてあるけど、十分足りた。

他にもいろいろな解説が生えてて、ユーザ解説 by ripity なんかは最初の解法と同じ部品でもっと速そう。この解説で求めている suffix の長さが、Bonus : オンラインで解いてください by toam のlcaの深さと同じもので、根付き木を実際に組み立てなくてもいいよってことなのかな。

ripity=サン方式

やれそうなのでやってみる。

import Data.Array

abc379f :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc379f _n _q hs = map f
  where
    rqa = makeRQArray min maxBound $ map fst $ scanr step (0, []) hs
    step h = final . until p g
      where
        p (cnt,hs) = cnt == 0 || h <= head hs
        g (cnt,hs) = (pred cnt, tail hs)
        final (cnt,hs) = (succ cnt, h:hs)
    f (l:r:_) = pred $ queryRQArray rqa l r

-- 不変セグメント木は省略、1つめの解と同じ

ここにあるコードの中では 565ms, 135MB で最も性能がよいものになった。

G - Count Grid 3-coloring

問題 ABC379G

シグネチャを決める。

abc379g :: Int      -- H
        -> Int      -- W
        -> [String] -- Sij
        -> Int      -- 答え

最初 $H,W \leq 200$ と読み間違えて、DPの状態数が途方もなくて無理、とか思ってしまった。
$H \times W \leq 200$ なので、短辺の長さの最大値は 14 $(14^2 = 196)$ となる。

考える

アライさんの助言に従って、$H$ と $W$ の小さい方をDPの状態を作るために使う。
縦長にしておくというか。

abc379g h w ss
  | h < w     = solve $ transpose ss
  | otherwise = solve ss

解説を見ると、ラスタースキャンで考えて、直前 W マスの割り当て状況をキー、
その場合の数を値としてDPの状態を構築して、1マスずつ進めていくものらしい。

そうとは知らずに、行ごとに、横 W 文字の内容をキー、値を場合の数として、H 段のDPをすると考えた。

すると、文字列をキー、場合の数の整数を値とする Map では重そうなので、文字列の内容を1文字2ビットに圧縮した 28ビットの整数をキーにして IntMap で管理するか。
しかしこうすると、現在の行のある埋め方に対して、前の行で違反していない全ての項目を列挙するのに、全てを舐めるので大変な計算量になりそう。
特に、左の方の文字で重複が確定してもそれより右の文字について全て読み飛ばすなど、無駄が大きい気がする。

そこでより効率的な持ち方として、文字1,2,3からなるTrieをつくり、末端に場合の数を入れることを考える。

data Tree = Empty | Leaf Int | Node Tree Tree Tree | Zeroth

Emptyは枝がない、Leafは末端、Nodeは文字 1,2,3 に応じた部分木が伸びる分岐。
Zerothは第0行で、どんな列とも干渉しないという特別な木とする。

DPとしては、注目する行に対応する木を文字列を '?' の自由度に応じて全て作り、
末端には、そこに至る文字列に対して、
直前の行で干渉しない全ての末端を探し、その場合の数の総和を入れる、を繰り返す。

最終行の木について、末端の値の合計が答えとなる。

この手順をコードにするとこうなる。

solve :: [String] -> Int
solve ss = finalize $ foldl' nextTree Zeroth ss

わかりやすいところから、最後の合計をとる関数。
acc の結果を先行評価することでspace leakを回避したつもり。

finalize :: Tree -> Int
finalize t = loop t 0
  where
    loop Empty acc = acc
    loop (Leaf x) acc = add acc x
    loop (Node t1 t2 t3) acc = loop t3 $! loop t2 $! loop t1 acc

nextTree の部品として、確定したひとつの列に対して、前の行の干渉しない全ての末端を合計するところ。
行の注目する文字に対して、それと異なる文字の枝2本について再帰すればよい。普通のTrieと逆。

これは上の finalize の延長上にあって、逆にもう少しロジックを追加すればそれを兼ねるようにもできる。

accumTreeFor :: Tree -> String -> Int
accumTreeFor Zeroth _ = 1
accumTreeFor t s = loop t s 0
  where
    loop Empty _ acc = acc
    loop (Leaf x) _ acc = add x acc -- s は ""
    loop (Node t1 t2 t3) ('1':cs) acc = loop t2 cs $! loop t3 cs acc
    loop (Node t1 t2 t3) ('2':cs) acc = loop t1 cs $! loop t3 cs acc
    loop (Node t1 t2 t3) ('3':cs) acc = loop t1 cs $! loop t2 cs acc
    loop (Node t1 t2 t3) ('*':cs) acc = loop t1 cs $! loop t2 cs $! loop t3 acc

finalize' t = accumTreeFor t $ repeat '*'

もっと頑張ると、部分木の値を使うか、使わずにバイパスするか、という関数を被せて場合分けをまとめることもできた。

accumTreeFor :: Tree -> String -> Int
accumTreeFor Zeroth _ = 1
accumTreeFor t s = loop t s 0
  where
    loop Empty _ acc = acc
    loop (Leaf x) _ acc = add x acc -- s は ""
    loop (Node t1 t2 t3) (c:cs) acc =
        byp '3' (loop t3 cs) $! byp '2' (loop t2 cs) $! byp '1' (loop t1 cs) acc
      where
        byp d v = if d == c then id else v -- バイパス

nextTree の中心のロジック、現在の行に対する木を構築する計算を考える。
直前に出力した文字 $p$ と、次に出力しようとする文字 $x$ によって場合分けする。
場合分けは並立でなく上優先。

• $p = x$ のとき：横に干渉しているのでこの枝は刈る。Empty を返す。
• （$p \neq x$ でかつ）$x \in \texttt{"123"}$ のとき：必ずその文字にするだけなので、その文字の枝だけ伸ばし、残り2つは Empty にする。
• （$x = \texttt{'?'}$ でかつ）$p \in \texttt{"123"}$ のとき：$p$ と異なる2文字の枝だけ伸ばし、$p$の枝は Empty とする。
• （$x = \texttt{'?'}$ かつ $p = \texttt{'*'}$ であるような）それ以外のとき：これは行頭のとき。横方向に縛りはないので、3本の枝を全て伸ばす。

コードにするとこうなる。現在位置が実際どんな文字列になったかを（逆順で）累積しておき（cs）、行の文字列が尽きたところで accumTreeFor を使って貼り付ける値を求める。
$p$は head cs なのでまとめる。先頭の場合の目印 '*' は null cs に読み替える。

nextTree :: Tree -> String -> Tree
nextTree t s = loop "" s
  where
    loop cs "" = Leaf $ accumTreeFor t $ reverse cs
    loop cs (x:xs) =
      case (cs, x) of
        (p:_, _) | p == x -> Empty
        (_, '1') -> prune $ Node t1    Empty Empty -- 指定があるときは従う
        (_, '2') -> prune $ Node Empty t2    Empty
        (_, '3') -> prune $ Node Empty Empty t3
        ('1':_, _) ->       Node Empty t2    t3    -- '?'のときは前と違う文字両方
        ('2':_, _) ->       Node t1    Empty t3
        ('3':_, _) ->       Node t1    t2    Empty
        _          ->       Node t1    t2    t3    -- 行頭の'?'は全部
      where
        t1 = loop ('1':cs) xs
        t2 = loop ('2':cs) xs
        t3 = loop ('3':cs) xs
    prune (Node Empty Empty Empty) = Empty
    prune t = t

再帰的に枝を伸ばしたとき、入力からの指定のため、途切れることがある。
選択肢が複数ある場合はどちらかが必ず生き残るが、ひとつだけのときに Node Empty Empty Empty という少し無駄な構造ができるのを、prune で剪定している。

accumTreeFor のときのように、バイパス関数を使って3とおりの場合分けを一つにまとめることもできるが、
バイパス関数も3とおり必要になり、コードの声が聞こえにくくなったのであえてこのままにしておく。

モジュロ演算の定義は省略。これで完成。
AC:2552ms, 39MB

公式解説のやり方

一マスずつDPを進める方法も気になるけど、いつかやる箱に入れて今日はここまで。