ABC329 A～F をHaskellで

A - Spread

問題 ABC329A

シグネチャを決める。

abc329a :: String  -- S
        -> String  -- 答え

間に空白を挟む、と言葉で言うと簡単だけど、
「要素の後に空白を入れる」だと最後に余計な空白が出てしまうので、
「一つ以上出力済み」というフラグを用意して、このフラグが立っているときだけ、要素の前に空白を入れる、とやる必要がある。

けれどもそういう面倒を引き受けてくれる便利な関数が Data.List にある。

intersperse ::  a  ->  [a]  -> [a]
intercalate :: [a] -> [[a]] -> [a]

結果

abc329a = intersperse ' '

実は、AtCoderのジャッジは優しくて、末尾に余計な空白があっても許される。

abc329a = concatMap (\c -> c : " ")

B- Next

問題 ABC329B

シグネチャを決める。

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

重複を除いてソートをして大きい方から二つめを取り出せばいいのだけど、計算量が $O(N \log N)$もかかって無駄なので、もっと節約することを考えると、

• 最大値を探し、それを除いたもののなかの最大値を探す
• 最大値と2番目を同時に探す、そのとき重複する値は捨てる

というやり方ならどちらも$O(N)$でできる。

結果

abc329b _n as = maximum $ filter (amax /=) as
  where
    amax = maximum as

abc329b _n as = loop minBound minBound as
  where
    loop _  a2 [] = a2
    loop a1 a2 (a:as)
      case compare a a1 of
        GT -> loop a  a1 as
        EQ -> loop a1 a2 as
        LT -> loop a1 (max a a2) as

C - Count xxx

問題 ABC329C

シグネチャを決める。

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

一つの文字種、例えば X について、$S$ のどこかに XXX があり、また別のどこかに XXXXX があるとき、X の長さ1から長さ5の5通りの部分文字列が作れる。
つまり、文字種ごとに、それが連続する最長の長さが、その文字種で作れる部分文字列の通り数になるので、これを足し合わせればよい。

結果

import Data.Array

abc329c :: Int -> String -> Int
abc329c _n =
  sum . elems .                        -- 4.足し合わせる
  accumArray max 0 ('a','z') .         -- 3.配列に最大長さを蓄積して
  map (\xs -> (head xs, length xs)) .  -- 2.文字とその長さの対にして
  group                                -- 1.同じ文字の並びごとに区切り

D - Election Quick Report

問題 ABC329D

シグネチャを決める。

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

候補者→得票数 という配列だけでは、当選者を決めるのに毎回$O(M)$かかってしまう。
得票数→候補者リスト という写像にすると、開票するたびに更新が大変な手間になる。
これが両方ともあれば、当選者は後者の写像の最大キーの最小キーとして即座に取り出せて、開票で更新するべき箇所は前者を使うことですぐわかる。

結果

候補者→得票数配列も、immutableなIntMapで代用する。
マイナス票はないので、候補者集合が空集合になっている票数に人を探しに行く心配はない。
loopの第2引数の初期値は 「0票→全員」IM.singleton 0 (IS.fromList [1..n]) が正しいが、ここが検索されることはないのでこのままで問題ない。

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.List

abc329d :: Int -> Int -> [Int] -> [Int]
abc329d _n _m = loop IM.empty IM.empty

loop :: IM.IntMap Int       -- 候補者→得票数
     -> IM.IntMap IS.IntSet -- 得票数→候補者集合
     -> [Int]               -- 開票
     -> [Int]               -- 当選者
loop _   _   []     = []
loop i2c c2i (a:as) = res : loop i2c1 c2i1 as
  where
    i2c1 = modMap succ a i2c
    cnt1 = i2c1 IM.! a
    c2i1 = modMap (IS.insert a) cnt1 $
           modMap (IS.delete a) (pred cnt1) c2i
    res  = IS.findMin $ snd $ IM.findMax c2i1

-- 空かもしれないIntMapの要素を変更する
modMap :: a               -- デフォルト値（空のとき、これが入っているものとする）
       -> Int             -- キー
       -> (a -> a)        -- 変更
       -> IM.IntMap a     -- マップ
       -> IM.IntMap a
modMap d i f im = IM.insert i (f $ IM.findWithDefault d i im) im

公式解説

公式解説によると：
直前までの開票での、当選者と、その得票数を保持する。
もう1票開票したとき、状態の変化は、その得票者が頭ひとつ抜けて唯一のトップ得票者になるか、追いついてトップ得票メンバーに加わるか、まだ追いつけない（このとき当選者は変化なし）、のいずれかである。

なるほど、全ての得票数についてメンバーを持っておく必要は全くなくて、トップのところだけ、しかも最小の候補でよかったのね。せっかくなので今回は落ち着いて mapAccumL で書くことにする。

abc329d :: Int -> Int -> [Int] -> [Int]
abc329d _n _m as = snd $ mapAccumL step (IM.empty, -1, 0) as

type State = ( IM.IntMap Int  -- 得票数カウント
             , Int            -- 当選者（トップ得票者の最小ナンバー）
             , Int)           -- トップ得票数

step :: State -> Int -> (State, Int)
step (cntmap, winner, topcnt) a =
  case compare cnt topcnt of
    GT -> ((cntmap1, a      , cnt   ), a)        -- aが頭一つ飛び出した
    EQ -> ((cntmap1, winner1, topcnt), winner1)  -- aがトップに追いついた
    LT -> ((cntmap1, winner , topcnt), winner )  -- aはまだトップに追いつかない
  where
    cnt = succ $ IM.findWithDefault 0 a cntmap
    cntmap1 = IM.insert a cnt cntmap
    winner1 = min a winner

なるほどこれはすっきり。
最初は必ず誰かが飛び抜けるので、開票数0での当選者はundefinedで構わない。

追記

状態に出力が完全に含まれているので、同じロジックでもっと短くできた。

abc329d :: Int -> Int -> [Int] -> [Int]
abc329d _n _m as = map snd $ tail $ scanl step ((IM.empty, 0), -1) as

type State = ( ( IM.IntMap Int  -- 得票数カウント
               , Int)           -- トップ得票数
             , Int)             -- 当選者（トップ得票者の最小ナンバー）

step :: State -> Int -> State
step ((cntmap, topcnt), winner) a =
  case compare cnt topcnt of
    GT -> ((cntmap1, cnt   ), a)             -- aが頭一つ飛び出した
    EQ -> ((cntmap1, topcnt), min a winner)  -- aがトップに追いついた
    LT -> ((cntmap1, topcnt), winner )       -- aはまだトップに追いつかない
  where
    cnt = succ $ IM.findWithDefault 0 a cntmap
    cntmap1 = IM.insert a cnt cntmap

E - Stamp

問題 ABC329E

シグネチャを決める。

abc329e :: Int     -- N
        -> Int     -- M
        -> String  -- S
        -> String  -- T
        -> Bool    -- 答え

$X$に、$T$の書かれたシールをうまく重ねて貼り付けることで、$S$と同じ文字列が作れるかと聞かれている。

時計を戻して考えると、$S$から始めて、重なりがなく完全な $T$ のシールが見えているものを剥がすと、その下に書かれていた文字はシールにくっついて読めなくなっている。これはもう # だと見なす。
次から、$T$ が完全に見えているか、一部が # で見えなくなっているシールは剥がしてよいとして、次々に剥がしていき、全てを # にできれば成功。

前回のDを彷彿とさせる。あちらはABCを消して切り詰めたが、こちらは # に置き換えられる。
#はワイルドカードなので、効率的な貼り順（剥がし順）でなくとも、貪欲にやっていっても結果は同じ。

左から順に見ていく。
ある箇所について剥がすことができたとき、その範囲は全てワイルドカードになる。初めてワイルドカードに変化した最も左の文字が、次に剥がすことのできる可能性のある$T$の最も右の文字である可能性があるので、そこまで戻して考え直す。（先頭まで戻らなくてよいということ。）
細かいのと、しょせん$M \leq 5$なので、固定で$|M|$文字戻してしまっても大差ない。
計算量的には、シール剥がし1回で1文字は消えるので、可能な回数はたかだか$N$回である。
右まで到達したとき、全てが # になっていれば成功。
几帳面なら、途中で変更した文字数を数えてもいいし、後で調べてもいい。

結果

abc329e _ m s t = loop "" s
  where
    eq '#' _ = True
    eq c   d = c == d
    isOK = all ('#' ==)

    loop xs ys
      | ende         = isOK ys && isOK xs
      | cond1, cond2 = loop xs1 ys3
      | otherwise    = loop (head ys : xs) (tail ys)
      where
        (ys1, ys2) = splitAt m ys           -- m文字取ってみて
        ende       = length ys1 < m         -- 取れなければ終わり
        cond1      = any ('#' /=) ys1       -- 取ったm文字は#ばかりではなく
        cond2      = and $ zipWith eq ys1 t -- 取ったm文字がTと一致
        (xs2, xs1) = splitAt (pred m) xs    -- xsからM-1文字右に戻す
        ys3        = reverse xs2 ++ replicate m '#' ++ ys2 -- #に書き換えて巻き戻し

公式解説について

最初は自分も、書き換えた位置から派生した、さらに書き換えられる位置をキューに投入するやり方で書いたけれど、M-1までの巻き戻しで十分だと気づいてこうなった。

F - Colored Ball

問題 ABC329F

シグネチャを決める。$a,b$は手抜きする。

abc329f :: Int     -- N
        -> Int     -- Q
        -> [Int]   -- Ci
        -> [[Int]] -- query_i
        -> [Int]   --  答え

素直に、箱の並びをmutable arrayで、箱の内容を整数集合で表すことを考えてみる。
ここで、

	union	size
Set	$O\bigl(m \log\bigl(\frac{n}{m} + 1 \bigr)\bigr), \; 0 < m \leq n$	$O(1)$
IntSet	$O(n+m)$	$O(n)$

であることに注意する。係数はおそらく IntSet の方がずっと小さいが、オーダーの観点で、IntSet の線形オーダーは危険である。
なのでボールの番号は贅沢にも Data.Set.Set Int に入れることにする。

結果

ループ間で共有される可変状態が配列だけだったので、forM を使うことができ、計算した順に結果を繋いだリストを返すことができた。
ループ間で共有される状態がmutableな情報だけでない場合、foldMでそれを持ち回すことになり、結果は一度配列などに詰めておいて、全て終わってから返す、ような形にせざるを得なくなる。
（Control.Scanl.scanM があればいいのかな？）

import Control.Monad
import qualified Data.Set as S
import Data.Array.ST
import Control.Monad.ST

abc329f :: Int -> Int -> [Int] -> [[Int]] -> [Int]
abc329f n _q cs abs = runST $
  do
    arr <- newListArray (1,n) $ map S.singleton cs :: ST s (STArray s Int (S.Set Int))
    forM abs (\(a:b:_) -> do
      sa <- readArray arr a
      sb <- readArray arr b
      let sb1 = S.union sa sb
      writeArray arr a S.empty
      writeArray arr b sb1
      return $ S.size sb1 
      )

ちなみに、IntSetを使うとちゃんとTLEする。

公式解説について

「ボールを移す」が一つずつ行う必要があり、集合に対して破壊的に行われるとする。
入れるボールの数が$A$個、入れられる方の集合の要素数が$B$個のとき、$\log(B) + \log(B+1) + ...\log(B+A-1) = O(A \log(A+B))$ かかる。

問題文の言うとおりに素直にボールの移動を行うと、$A$に大きい方ばかりを指定する意地悪なテストケースで $O(NQ\log N)$となり間に合わない。
そこで、それぞれの大きさを気にして、$A$が小さくなるように大きい方の箱にボールを移し、移動が指示と逆なら後で箱の方を交換することで、計算量が抑えられる。

という話なのね。
Data.Set.unionの計算量を見ると、$m \leq n$というのはつまりその配慮は済んでいるので、気にする必要がなかった。

G - Delivery on Tree

木は結局、一度走査するだけしかできない。

根と、スタートとゴールの位置関係によって、対応が変わる。

• スタートからゴールまで、降りていく一方ならば、その経路から道草をくう場合、それらの部分木に余計な荷物が1増える。
• スタートからゴールまで、登る一方ならば、その経路から道草をくう場合、それらの部分木に余計な荷物が1増える。
• スタートから登り、途中で折り返して後は下る場合、折り返し地点の二つの子の訪問順序がその向きに固定され、登り経路と下り経路の両方について、同様に道草ペナルティを考える。

辺を下る時刻と上る時刻を考えて木を一直線に展開して、確実にボールを運ぶ区間について、辺に重みを加えていく。
さらに、二つの子の訪問順序が固定されていないノードについて、その両方の可能性について、道草ペナルティがどれだけになるかを数える。
ペナルティ含めて上限Kを超えないような経路の個数を数える。

というところまでは読めたけれど、最後のところで、場合をどうやって数えたらいいのか訳がわからない。
（あと、それ以前のところもちゃんと実装するのがとても面倒そう。）

なのでまた後日気が向いたらということで。