ABC307 A～G をHaskellで

書きかけを発掘したので完成させた。

A - Weekly Records

問題 ABC307A

シグネチャを決める。

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

Data.List.Split.chunksOf で7個おきに区切って和をとる。

結果

abc307a _n = map sum . chunksOf 7

-- 開催当時はなかったので自作
chunksOf _ [] = []
chunksOf k xs = let (as,bs) = splitAt k xs of as : chunksOf k bs

B - racecar

問題 ABC307B

シグネチャを決める。

abc307b :: Int       -- N
        -> [String]  -- Si
        -> Bool      -- 答え

$N \leq 100$ なので総当たり $O(N^2)$ する。
例2のような場合に騙されないように、同じものを取り出してこないことだけ気をつける。
同じ物を取り出さない総当たりは tails で作れる。

import Data.List

abc307b :: Int -> [String] -> Bool
abc307b _n ss = or
  [ reverse z == z
  | x:ys <- tails ss, y <- ys
  , z <- [x ++ y, y ++ x]
  ]

ByteString を使い、(++) も reverse も使わずに回文判定することで 2msを達成できた。

C - Ideal Sheet

問題 ABC307C

間違いなく説明するとこんな長大な文章になってしまって仰々しいが、白黒ビットマップが2枚あって、黒を真として、AとBをORで重ね合わせてXが作れるかを判断せよ、平行移動のみ裏返しとか回転とかなしで、という要求。

シグネチャを決める。
ビットマップの情報に名前を付けて節約する。

data Bitmap = Bitmap { h ::Int, w :: Int, b :: [String] }

abc307c :: Bitmap  -- シートA
        -> Bitmap  -- シートB
        -> Bitmap  -- シートX
        -> Bool      -- 答え

縦横10ドットまでなので、ビットフィールドで表しても32ビット整数で間に合う。といっても面倒なので黒ドット座標の集合で扱う。（ぜいたく）

AやBを、いい位置に位置合わせする必要がある。
その黒ドットのいずれかを、目標であるXのいずれかの黒ドットに位置合わせしてみて、その結果全ての黒ドットがXのいずれかと重なる、つまりXで白であるべき位置に黒ドットがはみ出ないようなずらし方を全て作る。
（黒ドットの個数が上限$10^2$なので$10^4$通り。）

AとBのずらし方それぞれについて総当たりで、重ねた結果、Xの黒ドットでAもBもカバーしない、取り残されたドットがないことを確認する。その組み合わせが正解。
（上の通りの二乗で$10^8$通り。ぎりぎり。）

結果

import qualified Data.Set as S

abc307c :: Bitmap -> Bitmap -> Bitmap -> Bool
abc307c sheetA sheetB sheetX = or [pX == S.union a b | a <- pAs, b <- pBs]
  where
    pX = mkSet (b sheetX)
    pAs = shift pX $ mkSet (b sheetA)
    pBs = shift pX $ mkSet (b sheetB)

type Points = S.Set (Int,Int)

-- 点集合を作る
mkSet :: [String] -> Points
mkSet xs = S.fromList [(i,j) | (i,s) <- zip [0..] xs, (j,'#') <-zip [0..] s]

-- shift pX pA は、pAをずらしてpXに漏れなく重なるものを全て作る
shift :: Points -> Points -> [Points]
shift pX pA =
  [ pA1
  | (x,y) <- S.elems pX
  , (a,b) <- S.elems pA
  , let (d,e) = (x-a, y-b)
  , let pA1 = S.map (\(a,b) -> (a+d, b+e)) pA
  , S.isSubsetOf pA1 pX
  ]

レコード型まで作って持ってきたのに、HもWも使わなかった。（いつものこと。）

余談

そもそも今さらABC307を見たのはとあるマシュマロがきっかけ。

さて、どっちとも書いてないので「簡単すぎる」のか「難しすぎる」のかわからないけど、ABC406Cが実装も考察も大変という話同様で、ビットマップだと思って動かしたり、いい感じにやろうとすると命令型言語ではややこしくなっちゃう、という話だったのかな？と。

D - Mismatched Parentheses

問題 ABC307D

シグネチャを決める。
Sが長いのでByteStringで受け取る。
（出力はストリームになるのでStringで。）

import qualified Data.ByteString.Char8 as BS

abc307d :: Int            -- N
        -> BS.ByteString  -- S
        -> String         -- 答え

前から見ていって、

• 開き括弧に出会ったときは、その位置をスタックに積む
• 閉じ括弧に出会ったときは、スタックトップの開き括弧と対応させる。
  スタックが空なら対応する相手はないのでスルー。

で対応を発見できる。
次に、括弧内を除いた文字列を作るために、開き括弧の文字へ、対応する閉じ括弧の位置を注釈づける。
この注釈がない文字はそのまま出力し、注釈がある場合はそこまで飛ばす。

結果

import Data.Array

abc307d :: Int -> BS.ByteString -> String
abc307d n bs = loop2 0
  where
    jis = loop1 0 []
    arr = accumArray (flip const) 0 (0, pred n) jis

    loop1 :: Int   -- 現在位置
          -> [Int] -- 開き括弧の位置スタック
          -> [(Int,Int)] -- 対応のとれた位置の組
    loop1 i js
      | i == n = []
      | c == '(' = next (i:js)
      | c == ')', not (null js) = (head js, i) : next (tail js)
      | otherwise = next js
      where
        c = BS.index bs i
        next = loop1 (succ i)

    loop2 :: Int -- 現在位置
          -> String -- 答え
    loop2 i
      | i == n = ""
      | j /= 0 = loop2 (succ j)
      | otherwise = BS.index bs i : loop2 (succ i)
      where
        j = arr ! i

手書きの反復的計算なのがダサい。我ながら。

E - Distinct Adjacent

問題 ABC307E

シグネチャを決める。

abc307e :: Int  -- N
        -> Int  -- M
        -> Int  -- 答え

わからなかったのでフレンズさんの話を聞く。

ナイーブなDP

1人めに配る数を 0 に固定して考える。
1人めから $i$ 人めまで数を配り終えたとき、$i$ 人めが受け取った数が $x$ であるような場合の数を $C_i[x]$ とする。

1人めは固定なので $C_1[0] = 1, C_1[x \neq 0] = 0$
$1 < i \leq N-1$ 人めは、前の人と同じ値は受け取れないので $\displaystyle C_i[x] = \sum_{y \in \{ 0, \cdots, M-1 \}, y \neq x} C_{i-1}[y] = \sum_{y=0}^{M-1} C_{i-1}[y] - C_{i-1}[x]$
$N$ 人めは0だけ選択できない他は上と同じ漸化式でよい。

場合の数は全体では $\displaystyle \sum_{x=1}^m C_N[x]$ となる。
1人めに配る数を $M-1$ まで変えた結果も同じなのでこれの $M$ 倍が答え。
ひと工夫として $C_1[0] = M$ とするとこの計算を省ける。

$C[\cdot]$ をリストで実装することができる。

abc307e :: Int -> Int -> Int
abc307e n m = summ (tail cN)
  where
    c1 = m : replicate (pred m) 0
    cN = iterate step c1 !! pred n

    step c = map (reg . (acc -)) c
      where
        acc = summ c

modBase = 998244353 :: Int

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

例1,2は通るが、例3で固まる。

$C_i[\cdot]$ を求めるのに加算が $O(M)$ 回、それを $N$ 回するので $O(MN)$ となり、$M,N \leq 10^6$ という制限ではいかにも無理。

一人めと同じかそうでないか

よく観察すると、$C_i[x \neq 0]$ は全て常に同じ値になっていることがわかる。
性質としても、

• 0 の次は 0 は選べない
• 0 でない値 x の次は、0か、0 でも x でもない $M-2$ 通りから選べる

という関係にあるので、そのように上の漸化式を組み直す。
$i$ 人めが受け取った数が 0 である場合の数を $Z_i$, 0 でない場合の数を $D_i$ とする。後者は上で言うところの $M-1$ 個の $C_i[x \neq 0]$ と同じ値である。

$Z_1 = 1$ ($M$にする小技もある), $D_1 = 0$
$Z_i = (M-1) \cdot D_{i-1}, D_i = Z_{i-1} + (M-2) \cdot D_{i-1}$

最終結果は $(M-1) \cdot D_N$

abc307e :: Int -> Int -> Int
abc307e n m = mul m1 $ snd cN
  where
    c1 = (m, 0)
    cN = iterate step c1 !! pred n
    m1 = pred m
    m2 = pred m1

    step (z, d) = z1 `seq` d1 `seq` (z1, d1)
      where
        z1 = mul m1 d
        d1 = add z $ mul m2 d

add x y = reg (x + y)
mul x y = reg (x * y)

さらに

漸化式を行列で書き直してまとめると次のようになる。

\left (
\begin{array}{l}
Z_N \\
D_N
\end{array}
\right )
=
\left (
\begin{array}{ll}
0 & M-1 \\
1 & M-2
\end{array}
\right )^{N-1}
\left (
\begin{array}{l}
Z_1 \\
D_1 \\
\end{array}
\right)

行列のべき乗を $O(\log N)$ で行えばさらに高速化できる。

import Data.List

abc307e :: Int -> Int -> Int
abc307e n m = reg $ m1 * dN
  where
    m1 = pred m
    m2 = pred m1

    mat = powerish matmat (1,0,0,1) (0, m1, 1, m2) (pred n)
    (_zN, dN) = matvec mat (m, 0)

type Mat = (Int,Int,Int,Int)

matmat :: Mat -> Mat -> Mat
matmat (a,b,c,d) (e,f,g,h) = (reg $ a * e + b * g, reg $ a * f + b * h, reg $ c * e + d * g, reg $ c * f + d * h)

matvec :: Mat -> (Int,Int) -> (Int,Int)
matvec (a,b,c,d) (x,y) = (reg $ a * x + b * y, reg $ c * x + d * y)

powerish :: (a->a->a) -> a -> a -> Int -> a
powerish mul i a b = foldl' mul i [p | (True, p) <- zip bs ps]
  where
    bs = map odd $ takeWhile (0 <) $ iterate (flip div 2) b
    ps = iterate (\x -> mul x x) a

F - Virus 2

問題 ABC307F

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

abc307f :: [Int]   -- N,M,K,D
        -> [[Int,Int,Int]]  -- Ui,Vi,Wi
        -> [Int]   -- Ai
        -> [Int]   -- Xi
        -> [Int]   -- 答え
abc307f [n,m,k,d] uvws as xs = ...

ある意味ストレートなシミュレーション

ダイクストラ法の変種を作ることである程度はいけそう。

• 出発地点だけからエージェントが出発するところを、初期感染者のいる全ての部屋から出発させる
• 単純な経過時間の代わりに、現在何日目で、当日中の移動時間の残量というペアを使う
• エージェントがノードに到達すると、そこから生える辺について接続先のノードへの到着時刻を計算するが、
  単純な「到着時刻＋辺の長さ」の代わりに、
  • 当日中の移動時間残量が足りるならそれで移動
  • 足りない場合は翌日以降に持ち越すので、$W_i$ 以上を移動できる直近の日に移動開始

移動できる直近の日

これを特定するのに、辺ごとに$X_i$のリストを線形探索していては計算量ドブである。
使う場面から考えると、$W_i$ 以上という条件で一度だけ検索をし、その結果としてそのような全ての日付 $j$ の一覧、もしくは日付とその日の移動量 $X_j$ を取り出し、これを現在の日付超という条件で一度だけ検索し、最小の $j$ を見つけたい。

つまり、移動時間 $X_k$ をキーとし、「それ以上の $X_j$ をもつ全ての $j$ の集合」をひくマップがあるとうれしい。

例えば、$X_k$ を座標圧縮して添字にし、$j$ の集合を値、和集合を演算とするセグメント木を作ることが考えられる。
しかしこれは毎回の問いあわせで $O(\log D)$ 回の和集合演算が発生するのがつらいし、$2D$ 個の集合をメモリ上に保持するのも多め。
セグメント木は更新も意図したデータ構造なのに更新はなく問いあわせだけなのも間違いくさい。

更新が不要な状況で効率的な問いあわせだけを達成するのは累積和。つまり、

• 移動距離 $X_j$ をキーとし、ちょうどその移動時間をもつ日付 $j, j', j'', \dots$ だけの集合をひくマップを作る
• $X_j$ の大きい方から和集合演算で累積和をとり、$X_j$ 以上の移動時間をもつ全ての日付の集合を全て作っておく

とすればよい。本来、累積和で効率的なクエリとは、さらに差分をとることまでセットになっているが、この問題ではそもそも累積和の値が欲しかったのでちょうどいい。

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

    (_, x2j) = IM.mapAccumRWithKey x2jf IS.empty $             -- キーの大きい方から和集合で累積
               IM.fromListWith IS.union $                      -- マップを作り
               [(x, IS.singleton j) | (x,j) <- zip xs [1 ..]]  -- キー Xj 値 {j}
    x2jf is1 _k is2 = let is3 = IS.union is1 is2 in (is3, is3)

-- Xiを取り出す配列
    x = listArray (1,d) xs

移動先に到着する時刻

優先度付きキューで組み込みの大小関係を使うために、ここでの時刻を次のように定義する。

type Time =
  ( Int   -- 日付 未確定のとき maxBound
  , Int ) -- 当日中の移動時間残量 符号反転させて扱う

エージェントがノードに到着するイベントを管理する優先度付きキューの要素は次のようになる。

import qualified Data.Heap as PQ

type Event = PQ.Entry
  Time -- 到着時刻
  Int  -- 到着した頂点の番号

エージェントが時刻 (da, ea) にノード u に到着したとき、そこから生える辺 (v, w) <- g ! u について、v に到着する予定時刻は x2j を用いて次のように求められる。

import Data.Maybe

    distrib :: Event -> [Event]
    distrib (PQ.Entry (da, ea) u) = catMaybes
      [ if eaw <= 0
        then Just $ PQ.Entry (da, eaw) v    -- 当日中に移動できるとき
        else do                             -- 翌日以降に持ち越すとき
          (_, js) <- IM.lookupGE w x2j      -- w以上移動できる日で
          db      <- IS.lookupGT da js      -- 今日daより後ろの直近の日
          return $ PQ.Entry (db, w - x ! db) v
      | (v, w) <- g ! u                      -- uから生えている辺 v, w について
      , let eaw = ea + w                     -- w経過後の（マイナス）残り時間
      ]

    g = accumArray (flip (:)) [] (1, n) $ concat -- いつものグラフ表現
        [[(u,(v,w)), (v,(u,w))] | u:v:w:_ <- uvws]

という改変を含むダイクストラ法

作り置きのものを直すにしても、理解していないとできない。

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

abc307f :: [Int] -> [[Int]] -> [Int] -> [Int] -> [Int]
abc307f [n, _m, _k, d] uvws as xs = runST $
  do
    dist <- newArray (1, n) (maxBound, 0) :: ST s (STArray s Int Time) -- ノードの到着時刻を記録する配列
    forM_ as (\a -> writeArray dist a (0,0))                           -- 初期感染者
    loop <- fixST $ \loop -> return $ \queue -> -- キューを消費追加して dist を更新するメインループ
      case PQ.uncons queue of
        Nothing -> return ()                        -- イベントキューが空なら仕事終わり
        Just (ev1@(PQ.Entry t1 u), queue1) -> do    -- 先頭は時刻t1にuへ到着
          du <- readArray dist u                    -- 到着したノード u の記録時刻
          if du < t1
          then loop queue1                          -- 遅刻したなら仕事なし
          else do
            queue2 <- foldM (\q ev@(PQ.Entry t v) -> do  -- uの隣接vに時刻tに到着する、必要なイベントをキューに追加
              dv <- readArray dist v                     -- 移動先ノード v の記録時刻
              if dv < t
              then return q                            -- 劣るエージェントは出発しない
              else do
                writeArray dist v t                    -- 記録を更新するエージェントを出発させる
                return $ PQ.insert ev q
              ) queue1 $ distrib ev1
            loop queue2
    loop q0

    map (post . fst) <$> getElems dist -- maxBoundを-1に置き換える後処理
  where
    q0 = PQ.fromList [PQ.Entry (0, 0) a | a <- as] -- 初期キュー：初期感染者に対応するエージェントが並ぶ

-- daがmaxBoundなら1にする
    post da = if da == maxBound then -1 else da

提出：3341ms, 500MiB 通ったからヨシ！

G - Approximate Equalization

問題 ABC307G

シグネチャを決める。

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

最終的に差が1以内になるようにする。ということは、
$\big (\sum A_i \big ) \div N = q \dots r$ として、$r$ 個を $q+1$ に、残り $N - r$ 個を $q$ にする他ない。
この平均値 $q$ を引いて正規化する。
$B_i = A_i - q$
と、このデコボコを均すコストを最小化することが目標になる。

$B > 0$ のとき、余計な分を全て、または1残して隣に押しつける。
$B < 0$ のとき、周辺から略奪して自分を0または1にする。
最終的にどこから持ってくるのかは知らないが、隣から持ってくるのには、持ってくる量だけのコストがかかる。

端 $B_1$ から順に、その枠を0にするか1にするかの2択を行い、1にした枠の個数に対して、その状態にする最小コストを求めるDPをする。
このとき、操作後の $B_i$ は、それまでの値の累積和と 1 にした枠の個数で一意に定まるので、DPの引数/添字にはしない。

結果

import Data.List
import Data.Array.Unboxed

abc307g :: Int -> [Int] -> Int
abc307g n as = arrN ! r
  where
    (q, r) = divMod (sum as) n
    bs = map (subtract q) as
    arrN = foldl' step arr0 $ scanl1 (+) bs
    arr0 :: UArray Int Int
    arr0 = listArray (0,0) [0]
    step arr acc = arr1
      where
        (_,ub) = bounds arr
        ub1 = min r $ succ ub
        arr1 = accumArray min maxBound (0, ub1) $ concat
          [  (     k, cost + abs (acc - k    )) :
            [(succ k, cost + abs (acc - k - 1)) | k < ub1]
          | (k, cost) <- assocs arr
          ]

DP配列の範囲を最初から $(0, r)$ とする代わりに、上限は $r$ に達するまで1ずつ増やす形にした。
下限も、最後は $r$ だけになるように途中から持ち上げることもできるのだが、多少無駄な計算があってもいいことにした。というかその仕掛けが逆に実行時間を消費する。

Ex - Marquee