ABC383 A～F をHaskellで

A - Humidifier 1

問題 ABC383A

加湿器なんだから、穴が空いている云々は余計なのでは。
ていうか、A問題から入力データが多いね。
それはともかくシグネチャを決める。

abc383a :: Int     -- N
        -> [[Int]] -- Ti, Vi
        -> Int     -- 答え

最後に水を入れた時刻と、その時点の水量を持っておく。
次に水を入れる時刻になったら、それまでに失われる量を引いてから（下限0に注意）さらに水を足す。

結果

abc383a n tvs = loop 0 0 tvs
  where
    loop _ u [] = u
    loop s u ([t,v]:tvs) = loop t u1 tvs
      where
        u1 = max 0 (u - (t - s)) + v

B - Humidifier 2

問題 ABC383B

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

abc383b :: [Int]    -- H,W,D
        -> [String] -- Sij
        -> Int      -- 答え

$H, W \leq 10$ なので床は最大 100 個、これを ${}_n C_2$ とおり総当たりすると $100 * 99 / 2 \fallingdotseq 5000$ とおり、毎回全ての床についてどちらかの加湿器で加湿されるか判定すると100回、合わせて $5 \times 10^5$ 回の計算で力任せに。

結果

abc383b [h,w,d] ss = maximum [cnt h1 h2 | h1:h2s <- tails ijs, h2 <- h2s] -- 総当たりで
  where
    ijs = [(i,j) | (i,s) <- zip [1..] ss, (j,'.') <- zip [1..] s] -- 床マス
    mhd (a,b) (c,d) = abs (a - c) + abs (b - d)          -- マンハッタン距離
    cnt h1 h2 = length
      [() | p <- ijs, mhd h1 p <= d || mhd h2 p <= d]    -- h1かh2で加湿される床の個数

C - Humidifier 3

問題 ABC383C

シグネチャはBと同じ、っぽいけど、部屋が広いので ByteString を使おう。

import qualified Data.ByteString.Char8 as BS

abc383c :: [Int]           -- H,W,D
        -> [BS.ByteString] -- Sij
        -> Int             -- 答え

複数個の加湿器から同時にペンキが流れ出して、距離Dまで伸びたら止まる様子を計算する。
幅優先探索は、座標に加えて寿命も持たせたエージェントをFIFOキューに入れて管理するのがよくあるやり方。
自分の好きなやり方は、今回の世代で処理するべき座標のリストを消費しつつ、誕生した次の世代を別のリストにため込み、前者が空になったら世代交代する、このとき寿命をカウントダウンする、という構造。寿命が全員同じであることが使える条件だけど、それ普通。

速度のために、訪問済みのマスをmutable arrayで管理しよう。

結果

import Data.Array.Unboxed

abc383c :: [Int] -> [BS.ByteString] -> Int
abc383c [h,w,d] ss = runST $
  do
    bm <- newArray bnds False :: ST s (STUArray s (Int,Int) Bool)
    loop bm d ps0 []
    length . filter id <$> getElems bm
  where
    bnds = ((1,1),(h,w))
    grid = listArray bnds $ concatMap BS.unpack ss :: UArray (Int,Int) Char
    ps0 = [h | (h,'H') <- assocs grid]
    neighbors (i,j) =
      [ q
      | q <- [(pred i,j),(succ i,j),(i, pred j),(i, succ j)]
      , inRange bnds q
      , grid ! q /= '#'] -- bmで塗られていない、はimmutableには計算できないので放置

    loop bm e ps next
      | e < 0     = return ()
      | null ps   = loop bm (pred e) next []
      | otherwise =
      do
        let (p:ps1) = ps
        bmp <- readArray bm p
        if bmp then loop bm e ps1 next else do        -- 塗装済みなら仕事なし
          writeArray bm p True                        -- 新しい場所なら、塗って、
          loop bm e ps1 (neighbors p ++ next)         -- 隣接するマスを次のターンのために登録

D - 9 Divisors

問題 ABC383D

うってかわって引数は整数ひとつ。

abc383d :: Int  -- N
        -> Int  -- 答え

数を素因数分解して $x = p_1^{k_1} \times \dots \times p_n^{k_n}$ となるとき、その約数は、これらの素因数をそれぞれいくつ使うか、というバリエーションの個数になるので、0個の場合を足して $\prod (k_i + 1)$ となる。
約数が9個になる数は、素因数の種類 $(n)$ がいくつでそれぞれの個数 $k_i$ がいくつになればいいのか考えてみると、

• $n=1$ のとき $k_1 = 8$
• $n=2$ のとき $3 \times 3 = 9$ なので $k_1 = k_2 = 2$
• $n \geq 3$ のとき、9にできない

$n=1$ の場合は、8乗して$N$以下になるような素数の8乗がそれなので、それを数える。
$n=2$ の場合は、素数の平方を考えて、それらの異なる二つを掛けた結果が$N$以下である組み合わせの個数。
$p_1 < p_2$ とする。
調べる素数の上限 $q$ は、最小の素数 2 の場合に $2^2 \cdot q^2 \leq N$ を満たす最大の $q$ までとわかる。
ある $p_1$ に対して $p_1 ^2 \cdot p_2 ^ 2 \leq N$ で $p_2$ の上限が決まる。
$p_1$ が $i$ 番目の素数、上式の上限な $p_2$ が $j$ 番目の素数とすると、$j - i$ が組み合わせの個数。命令的なコードの場合、これが0または負になるとき、探索を終了できる。
命令的にしなくても、$p_1$ の上限は $p_1 ^4 \leq N$ で得られる。

$j$ を求めるのに、自分は、素数の平方に対してその番号を取り出す IntMap を使った。
公式解説によると、$i$ を増やすと $j$ は小さくなるので、尺取法の変種で線形探索できるという。

結果

$n = 1$ の場合については確定

abc383d :: Int -> Int
abc383d n = cnt1 + cnt2
  where
-- p^8 ≦ N であるような素数 p について p^8 は、p^0～p^8 の9個の約数を持つ
    cnt1 = length $ takeWhile ((n >=) . (^ 8)) primes

$n = 2$ の場合について、背番号を取り出す IntMap を使う方法：

    -- 素数の2乗に背番号
    im = IM.fromDistinctAscList $ zip (takeWhile (div n 4 >=) $ map (^ 2) primes) [1 ..]
    -- 4乗してNを越えない範囲でp1^2を振る
    p1is = zip (takeWhile ((n >=) . (^ 2)) $ map (^ 2) primes) [1 ..]
    -- p1^2 に対応する p2^2 の上限を探して個数を求める
    cnt2 = sum $ map cnt2f p1is
    cnt2f (p1,i) = max 0 $ j - i
      where
        Just (p2,j) = IM.lookupLE (div n p1) im

尺取法：

    -- 背番号に素数の2乗、p1側
    pis = zip [1 ..] $ map (^ 2) primes
    -- p2側は、上限までに切り詰めて逆転
    ips = reverse $ takeWhile ((div n 4 >=) . snd) pis
    -- 短くなる尺取法
    cnt2 = sum $ maomao pis ips
    maomao _ [] = []                            -- この終わり方もあるので注意
    maomao pipis@((i,p1):pis) ipips@((j,p2):ips)
      | i >= j       = []                       -- オーバーラン
      | p1 * p2 <= n = j - i : maomao pis ipips -- 見つけた。左を捨てて次のp1へ
      | otherwise    = maomao pipis ips         -- 見つけてない。p2が大きいので右を捨てる

アルゴリズム	時間	メモリ
IntMap	312ms	29MB
尺取法	281ms	25MB

$O(1)$の速さがわずかに結果に現れた。

E - Sum of Max Matching

問題 ABC383E

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

abc383e :: [Int]   -- N,M,K
        -> [[Int]] -- ui,vi,wi
        -> [Int]   -- Ai
        -> [Int]   -- Bi
        -> Int     -- 答え

まず、ひとつの経路のコストが、総距離ではなく、辺のコストの最大値であるところがポイント。
どんなに遠回りでも、安い辺だけを通って到達できるなら、その経路を使うということ。

この制約は、クラスカル法を連想させる。
安い方の辺から順に、何かを連結するような辺だけを選んで使っていくことで、最小全域木を見つけるアルゴリズムである。
「何かを連結する」かどうかは、Union-Find で判定する。

重さ $w$ の辺を追加したときに、分割 $X$ と分割 $Y$ が連結されるならば、
$X$ の中でまだ相手を見つけていない $A$ の要素と、$Y$ の中でまだ相手を見つけていない $B$ の要素について、
これらを対応させて $f(A_i,B_i) = w$ とするのが最安の経路である。
逆も同じ。

両者の未婚者数が同数でないときにあぶれるものがいて、誰をあぶれさせるか、を選択する部分問題が起きそうだが、実はそれらは互いに $w$ 以下のコストで行き来でき、また今後追加される辺の重さは $w$ 以上に決まっているので、誰を選んでも同じである。
さらにいうと、まだ未婚なのは誰か、を追跡する必要すらなく、$A$ の未婚者数と $B$ の未婚者数だけを把握しておけばよい。

さらにさらにいうと、ある分割の中に $A$ の未婚者と $B$ の未婚者が両方いるなら
「おまえらもうくっついちゃえよ」ということで、両方同時に存在することはありえない。
なので、正の数で $A$ の人数、負の数で $B$ の人数を表せば、整数ひとつを分割の属性として貼り付けるだけで済む。

以上をまとめると：

• 各頂点だけからなる単独の分割に、$A$ と $B$ の未婚者数を貼り付けておく。
• クラスカル法の流れで、必要な辺を洗い出していく
• 新たに連結が起きたとき、両者の未婚者数を合算して分割に貼り付け直す。
• またこのとき、結婚できたペアの数だけ、$w$ が総コストに算入される。
• 全員が結婚できた時点で止められるなら止めてもいいし、辺を使い切るまで計算を続けても最悪計算量は同じ

となる。

結果

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

abc383e :: [Int] -> [[Int]] -> [Int] -> [Int] -> Int
abc383e [n,_m,_k] uvws as bs = runST $
  do
    uf <- newUF (1,n) -- UnionFind
    mcnta <- thaw cnta :: ST s (STUArray s Int Int) -- 各代表ノードの分割が持つAの個数 - Bの個数
    foldM (\acc (u:v:w:_) -> do
      mab <- uniteUF uf u v
      case mab of
        Nothing -> return acc
        Just (a,b) -> do                     -- 連結が起きたら
          cntuu <- readArray mcnta a         -- 元の両方の分割のAB人数を
          cntvv <- readArray mcnta b
          writeArray mcnta b (cntuu + cntvv) -- 足せばよい
          return $ acc + w * marrige cntuu cntvv -- 結婚成立数×w だけ増やす
      ) 0 uvws1
  where
    uvws1 = sortBy (compare `on` (!! 2)) uvws
    cnta :: UArray Int Int
    cnta = accumArray (+) 0 (1,n) $ [(a,1) | a <- as] ++ [(b,-1) | b <- bs]
    -- 結婚成立数
    marrige c d
      | signum c == signum d = 0
      | otherwise = min (abs c) (abs d) -- 符号が異なるとき、絶対値の小さい方が居なくなる、その人数
    -- 別解
    marrage c d = div (abs c + abs d - abs (c + d)) 2 -- 人数の絶対値で減った数の半分

Union-Find は、ACL の移植版が導入されたら、それに合わせざるを得ないけど、それまでは自前で。
今までの自前のと比べて違う点は

• 要素数$N$に対して$(0,N-1)$で作る代わりに、下限と上限のペアを与える。つまり $(1,N)$ が指定できる
• 分割の要素数を維持しなくなった。
• 統合のありなしを、Bool で返していた代わりに、Maybe (Int,Int) で返すようにした。
  要素数やその他、分割の統合に際して何かしたいときは、このペアに元の代表元が渡されるのでそれを使えばできる。

-- データ表現
-- 自分の番号を指しているとき、自分が代表元
type UnionFind s = STUArray s Int Int

-- Union-Find構造体を作る
-- 番号の範囲を下限、上限で
newUF :: (Int,Int) -> ST s (UnionFind s)

-- 代表元を得る
getRoot :: UnionFind s -> Int -> ST s Int

-- 統合する。
-- 統合が実際に行われたとき、元の代表元2つをペアにして返す（sndが統合後の代表元）
uniteUF :: UnionFind s -> Int -> Int -> ST s (Maybe (Int,Int))

F - Diversity

問題 ABC383F

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

abc383f :: [Int]   -- N,X,K
        -> [[Int]] -- Pi,Ui,Ci
        -> Int     -- 答え

購入金額の上限 $X \leq 50000$ と控えめなのがヒントで、これがDP配列の添え字の一つに使えそう。
（というアプローチは、微妙にメタでインチキな気もするけど）

素朴に足し込んで行くと、今回の商品の色が初めてならばスコアを $U_i$ だけでなく追加で $K$ 増やす必要があるので、色のパターンの全ての組み合わせを添え字にするハメになりそう。
そうする代わりに、商品を色でグループ化して、ある色についてある期間考えたら、それ以降はその色について気にしなくてよくなるようにする。

色 $c$ まで考え終わったとき、合計購入価格 $x$ で達成できる最大の満足度を $score_c[0 \leq x \leq X]$ とする。
初期値は $score_0[0] = 0, score_0[x > 0] = -\infty$ とする。

次の色 $d$ のグループについて考えるとき、
「色 $d$ の商品を一つ以上購入していて、合計購入価格 $x$ で達成できる最大の満足度のmutable配列を $new_d[0 \leq x \leq X]$ とする。
初期値は $new_d[x] = -\infty$ とする。
それぞれの商品 $(P_i, U_i, C_i = d)$ について、全ての $0 \leq x \leq X - P_i$ について、

• $new_d[x] + U_i$ で $new_d[x + P_i]$ を更新する（大きい方に入れ替える）
  （mutable arrayでするときは、$x$の大きい方から順にする必要がある）
• $score_c[x] + U_i + K$ で $new_d[x + P_i]$ を更新する

色$d$について全て処理が済んだら、
$score_d[x] = \max(score_c[x], new_d[x])$
とする。

全ての色について処理が済んだら、
$\max_x score_N[x]$
が答え。

商品一つについて$new_d$の足し込みが$2X$回、商品$N$個なので計算量は $O(NX) = 25 \times 10^6$ なので間に合う。

結果

商品ごとの $new_d$ への足し込みはほぼ全域（$X-P_i$ 以下）に渡るので、immutable arrayでも大きな不利はない。
$K$を足し込むタイミングを、最初に購入した時点でなく、$score$に戻す時点に変更した。

import Data.Array.Unboxed

abc383f :: [Int] -> [[Int]] -> Int
abc383f [n,x,k] pucs = maximum $ elems scoZ
  where
    -- 色でグループ化
    puA :: Array Int [[Int]]
    puA = accumArray (flip (:)) [] (1,n) [(c, puc) | puc@(_:_:c:_) <- pucs]
    -- score_0
    sco0 :: UArray Int Int
    sco0 = listArray (0,x) $ 0 : repeat minBound
    -- 各 new_d の初期値
    newd0 :: UArray Int Int
    newd0 = listArray (0,x) $ repeat minBound
    -- score_* の最終値
    scoZ = foldl' cstep sco0 $ elems puA
-- 色ごとに入れてみるステップ
    cstep sco pus = accum max sco [(p, u + k) | (p,u) <- assocs newdZ, 0 <= u]
      where
        newdZ = foldl' istep newd0 pus
-- 色について要素を入れてみるステップ
        istep newd (p:u:_) = accum max newd
            [(p1, au + u) | (ap,au) <- assocs sco ++ assocs newd, 0 <= au, let p1 = ap + p, p1 <= x]

G - Bar Cover

「頑張るのだ！」というやつは大抵シャレにならない難易度なので撤退。