ABC328 A～G をHaskellで

- C : 累積和
- E, F : Union-Find

~~宿題：Fを重み付きUnion-Findで書く方法~~ 追記しました。

---

# A - Not Too Hard

[問題 ABC328A](https://atcoder.jp/contests/abc328/tasks/abc328_a)

シグネチャを決める。

```haskell
abc328a :: Int   -- N
        -> Int   -- X
        -> [Int] -- Si
        -> Int   -- 答え
abc328a _n x ss = sum $ filter (x >=) ss
```

問題文の言うとおりにした。

# B - 11/11

[問題 ABC328B](https://atcoder.jp/contests/abc328/tasks/abc328_b)

シグネチャを決める。

```haskell
abc328b :: Int   -- N
        -> [Int] -- Di
        -> Int   -- 答え
```

## 怠惰な解法

カレンダーの全ての日付を生成し、数字に直してみて、ゾロ目なものを数える。

```haskell
abc328b :: Int -> [Int] -> Int
abc328b n ds = length [() | (i, di) <- zip [1..] ds, d <- [1..di], allSame (show i ++ show d) ]

allSame (x:xs) = all (x ==) xs
```

## 考えすぎる解法

可能性のある月は 1～9 と、11の倍数 11～99 のみで、後者の場合、ゾロ目となる数字は 11 で割った商である。
数字 $d$ が決まると、可能性のある日は $d, 11d$ の二つで、$D_i$ を超えない範囲のものを数えればよい。

```haskell
abc328b :: Int -> [Int] -> Int
abc328b n ds = length [() | (i,di) <- zip [1..] ds, c <- cands i, c <= di]

cands :: Int -> [Int]
cands x
  | x <= 9 = [x, 11 * x]
  | r == 0 = [q, 11 * q]
  | otherwise = []
  where
    (q,r) = divMod x 11
```

# C - Consecutive

[問題 ABC328C](https://atcoder.jp/contests/abc328/tasks/abc328_c)

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

```haskell
abc328c :: Int     -- N
        -> Int     -- Q
        -> String  -- S
        -> [[Int]] -- li, ri
        -> [Int]   -- 答え
```

$S$の先頭から全ての$i$文字目までについて、同じ文字が隣り合う箇所の個数を数えて配列に入れておく。
あとは差を返すだけでよい。

## 結果

```haskell
abc328c n _q s lrs = [arr ! r - arr ! l | l:r:_ <- lrs]
  where
    arr = listArray (1,n) $ scanl (+) 0 $ zipWith f s $ tail s
    f c d = if c == d then 1 else 0
```

# D - Take ABC

[問題 ABC328D](https://atcoder.jp/contests/abc328/tasks/abc328_d)

シグネチャを決める。

```haskell
abc328d :: String  -- S
        -> String  -- 答え
```

スタックつきオートマトンのようなものを作ればうまくできる。
'C'に遭遇したときに、現在位置の手前に`B`,`A`がこの順にあったとき、それら全てを消費する。

## 結果

```haskell
abc328d :: String -> String
abc328d s = loop "" s
  where
    loop ('B':'A':t) ('C':s) = loop t s
    loop t (c:s) = loop (c:t) s
    loop t "" = reverse t
```

# E - Modulo MST

[問題 ABC328E](https://atcoder.jp/contests/abc328/tasks/abc328_e)

シグネチャを決める。$u_i, v_i, w_i$は手抜きする。

```haskell
abc328e :: Int     -- N
        -> Int     -- M
        -> Int     -- K
        -> [[Int]] -- ui,vi,wi
        -> Int     -- 答え
```

木は辺を頂点の個数-1だけ持つので、頂点$N$から始めて、頂点2まで順に辺を1本ずつ、既に連結になっていないものを選ぶ。
最後に、頂点1を含めて連結していれば全域木なのでそのコストを報告する。
最後に、それらの最小値をとる。

## 結果

Union-Findは今回、`uniteUF` が新規に連結を行ったとき更新値を、さもなくば`Nothing`を返す形にしている。

```haskell
abc328e :: Int -> Int -> Int -> [[Int]] -> Int
abc328e n _m k uvws = minimum $ loop newUF 0 n
  where
    g = accumArray (flip (:)) [] (1,n)                       -- いつものグラフ
        [p | u:v:w:_ <- uvws, p <- [(u,(v,w)),(v,(u,w))]]
    loop uf cost 1 = [mod cost k | snd (getRoot uf 1) == n]  -- 頂点1が全部でN個と連結なら総コストを返す
    loop uf cost a =
      [ d
      | (b,c) <- g ! a                      -- 頂点 a と隣接する頂点 b で
      , Just uf1 <- [uniteUF uf a b]        -- 既に到達可能でなかったものについて
      , d <- loop uf1 (cost + c) (pred a)   -- 再帰的にコストを求める
      ]
```

[提出](https://atcoder.jp/contests/abc328/submissions/47515602)

# F - Good Set Query

[問題 ABC328F](https://atcoder.jp/contests/abc328/tasks/abc328_f)

シグネチャを決める。$a_i, b_i, d_i$は手抜きする。

```haskell
abc328f :: Int     -- N
        -> Int     -- Q
        -> [[Int]] -- ai,bi,di
        -> [Int]   -- 答え
```

整数列 $(X_1, \dots, X_N)$ それぞれに頂点を対応させたグラフを考える。
問題の設定は、クエリ $(a, b, d)$ を前から順に、$X_a - X_b = d$ が成り立つものを集めると言っている。
このクエリは、$X_a$ と $X_b$ の差を指定していて、それが矛盾していれば採用しない。

とすると、$(a,b)$ を辺とすると、連結でない限り、矛盾が生じることはない。
既に連結されている者同士を次のクエリの辺が接続しようとするとき、矛盾があれば却下せよ、ということ。
なのでまず、初めて接続する辺だけからなる森を作る。
Union-Findで森の木ごとに代表点$r$を決め、$X_r = 0$ と固定し、連結する他の頂点の値を辺から木の探索で決める。

こうして$X_i$の値が全て決定されたら、それに矛盾しない$S_i$の$i$を全て出力する。

## immutableなHaskellに翻訳

Union-FindはE問題と同じ。

```haskell
abc328f :: Int -> Int -> [[Int]] -> [Int]
abc328f n _q abds =
-- Xa - Xb = d であるような i を順に挙げる
    [i | (i, a:b:d:_) <- zip [1..] abds, xm IM.! a - xm IM.! b == d]
  where
-- Union-Findで、スパニング木の森を作り、その辺の情報を取り出す
    (uf, es) = mapAccumL step newUF abds
    step uf abd@(a:b:_) =
      case uniteUF uf a b of
        Nothing  -> (uf , Nothing)
        Just uf1 -> (uf1, Just abd)

-- スパニング木の森なグラフ
    g = accumArray (flip (:)) [] (1,n)
        [p | Just (a:b:d:_) <- es, p <- [(a,(b,-d)),(b,(a,d))]]

-- それぞれの木の代表頂点
    roots = IS.fromList $ map (fst . getRoot uf) [1..n]

-- 代表はX=0とし、gをtraverseして他のXの値を全て求める
    xm = loop IM.empty [] (IS.elems roots)
    loop im [] [] = im
    loop im [] (r:rs) = loop im [(r,0)] rs
    loop im ((i,x):ixs) rs
      | IM.member i im = loop im  ixs  rs
      | otherwise      = loop im1 ixs1 rs
      where
        im1  = IM.insert i x im
        ixs1 = [(b, x + d) | (b,d) <- g ! i] ++ ixs
```

[TLE x 3](https://atcoder.jp/contests/abc328/submissions/47519875) で終わった。

## mutableな計算に翻訳

さすがにF問題に immutable な Union-Find は舐めプすぎたようなので、同じアルゴリズムを `IOVector` で再実装したら、 [385ms, 217MB](https://atcoder.jp/contests/abc328/submissions/47520799) でACした。

# G - Cut and Reorder

[問題 ABC328G](https://atcoder.jp/contests/abc328/tasks/abc328_g)

シグネチャを決める。

```haskell
abc328g :: Int    -- N
        -> Int    -- C
        -> [Int]  -- Ai
        -> [Int]  -- Bi
        -> [Int]  -- 答え
```

## 頭に入ってこない問題文

整理してみた。

長さ$N$の数列$A=(A[1],...,A[N]), B=(B[1],...,B[N])$が与えられる。
目的は、$A$に下記の操作を任意回行って$B$と等しくすることである。
それに必要な合計コストの最小値を求めよ。

- 操作1: 一つの要素の値を任意に変更する。  
整数$i$と整数$k$を任意に選び、$A'[i] = A[i] + k, A'[j] = A[j] \ (j \neq i)$ とした $A'$を新しい$A$とする。  
コストが $|k|$ かかる。
- 操作2: 数列を任意の断片に分割し、並べ替える。  
$X$個の断片に分割して並べ替えるとは、
分割の位置 $1 = i[1] < i[2] < \dots < i[X] < i[X+1] = N+1$ と$(1,2,...,X)$の順列$(p[1],\dots,p[X])$を自由に選び、
$A$の$X$個の断片 $AD[k] = (A[i[k]],A[i[k]+1],...,A(i[k+1]-1))$ $(1\leq k \leq X)$を$AD[p[1]], \dots, AD[P[X]]$ の順に連結したものを新しい$A$とする。  
コストが $(X-1)C$ かかる。

問題の内容は理解できた。
操作1は、切った箇所の個数でコストがかかるので、最初に一度だけ行えば十分である。$A$をいい感じに切って並び替えて、pointwiseに差がなるべく小さくなるようにして、差の合計＋切った箇所の個数×$C$を最小にすることが目的である。

そして、$N \leq 22$とか一見ずいぶん小さな問題に見えるところがとても曲者なのだろう。

<!--
以下で定義される数列$A'$に置き換える。

$A'[j] = \left \{ \begin{array}{ll} A[i] + k & i = j のとき\\ A[j] & それ以外のとき \end{array} \right .$
-->

## 休むに似たり

$N \leq 22$なので、$A$や$B$の使用済みの値を記憶する整数集合をビット表現で表したりするのだろうと（ズルい）考え方をして、次のようなDPをしてみた。

集めるDPでは表しにくそうなので、諦めて配るDPをする。
$B$を前から順に調べていく。
状態として、$A$の使用済みの要素の番号のビット配列と、直前の$B$を消すのに使った$A$の番号の対について、最小コストを記録しておく。
（巡回セールスマン問題のときに使うようなDP）
次に$B[j]$に（未使用の）$A[i]$を対応させるとき、コストは$|B[j] - A[i]|$と、直前に使った$A$が$i-1$でないときは追加で$C$かかる。
$B$を全て調べた後の、コストの最小値を答えとする。

```haskell
import Data.Bits
import qualified Data.Map as M

abc328g :: Int -> Int -> [Int] -> [Int] -> Int
abc328g n c as bs = minimum $ M.elems mN
  where
    mN :: M.Map (Int,Int) Int
    mN = foldl' step (M.fromList [((-1,0),0),((-2,0),-c)]) bs
    step m bi = M.fromListWith min
      [ ((j, setBit bmp j), cost + abs (aj - bi) + if succ i == j then 0 else c)
      | ((i, bmp), cost) <- M.assocs m
      , (j, aj) <- zip [0..] as
      , not $ testBit bmp j
      ]
```

小さな入力については正解するものの、$N=22$な例3についてはどうにもならなかった。

## フレンズさんの解説

[アライさんの解説](https://twitter.com/kyopro_friends/status/1723338259118580193)
$N$要素のビット配列で、$A$のうち使用済みの分を表す。そのビットの1の個数だけ、$B$は前から順に消費する。
そのような状況の最小コストをDPし、$B$の続きと対応させる$A$の区間を総当たりする。

ということでやってみる。
DP配列の代用として`IntMap`を使い、まず0にコスト0を入れた初期値を用意する。
とここで、ステップごとに連続する範囲を割り当てるたびに$C$を足すと、切れ目の数でなく断片の個数になってしまう。これを補填するため、初期値は$-C$にする。
そして、0から順に$2^N-2$までをより添え字の大きい要素に配り、配り終えた後の要素$2^N-1$が答え。
配り終えた要素はその後は不要なので捨ててやると、メモリの節約になるだろうか。

ここまでをコードにする。

```haskell
import Data.Bits
import qualified Data.IntMap as IM

abc328g :: Int -> Int -> [Int] -> [Int] -> Int
abc328g n c as bs = im IM.! all1
  where
    all1 = 2 ^ n - 1
    im = foldl' step (IM.singleton 0 (negate c)) [0 .. pred all1]
    step :: IM.IntMap Int -- ビットマップ→コスト 配列
         -> Int           -- 今回配る添え字 = Aの使用済み要素ビットマップ
         -> IM.IntMap Int -- 更新された配列
    step im bmp = IM.delete bmp im1 -- 最後に、配り終えた要素を削除（必須ではない）
      where
        im1 = ...                   -- bmpの要素を配り終えた更新済みの配列
```

コスト配列の注目要素を配るステップ動作`step`を考える。
添え字`bmp`の二進数での1の個数が`B`の前から使用済み要素数で、
今回はその次からと`A`のあらゆる（未使用の）箇所とを対応させる。
まずはその部分を取り出す。

```haskell
        bs1 = drop (popCount bmp) bs -- 今回対応を探すBの部分
```

配るコストの基本値は、配列の要素に$C$を足した値。

```haskell
        cost = im IM.! bmp + c -- 現状のコスト(今回のCを含む)
```

$A$の開始点は、添え字を$1$から$N$まで全て振り、未使用な要素全てを考える。
$A$のその位置から先と、$B$の差を順に求めておく。

```haskell
              [ ...
              | i <- [0 .. pred n], not $ testBit bmp i
              , let ds = zipWith (\b a -> abs $ b - a) bs1 $ drop i as
```

$A$の`i`から始まる断片について、未使用な範囲全ての長さについて、配る先と配るコストを計算する。
範囲を1ずつ広げて、`bmp`のビットを立てることを累積すると配る先のビットマップが作れる。
`cost`に`ds`の値を順に足し込むと、配るコストが作れる。
終わるのは、`bmp`のビットが1になったときである。

```haskell
              , let is = takeWhile (not . testBit bmp) [i .. pred n]  -- i以降でbmpが1になるまでの添え字の列
              , let bmps = scanl setBit bmp is                    -- 配る先のビットマップの列に直す
              , let cs   = scanl (+) cost ds                      -- dsを累積して配るコストの列を作る
```

この`bmps`の要素と`cs`の要素をそれぞれ組にして`im`に戻す。ただし先頭は今注目している位置なので捨てる。

```haskell
        im1 = foldl' (\im (b,c) -> IM.insertWith min b c im) im
              [ p
              | ...
              , p <- tail $ zip bmps cs
              ]
```

完成。例1,2は通るので正しそう。
しかしもちろん、$2^{22}$ 要素も抱え込む`IntMap`は時間も空間も間に合わない。

## 命令的な計算を導入

同じアルゴリズムを、mutable vectorで実装する。
上の`im1`のリスト内包表記を極力残す形で書いてみた。

```haskell
import Data.Bits
import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad

abc328g :: Int -> Int -> [Int] -> [Int] -> IO Int
abc328g n c as bs =
  do
    costs <- MUV.replicate (succ all1) maxBound
    MUV.write costs 0 (negate c)
    forM_ [0..pred all1] (\bmp -> do
      cost <- (c +) <$> MUV.read costs bmp
      forM_ (distrib bmp cost) (\(b1,c1) -> MUV.modify costs (min c1) b1)
      )
    MUV.read costs all1
  where
    all1 = 2 ^ n - 1 :: Int
    distrib bmp cost =
      [ p
      | let bs1 = drop (popCount bmp) bs
      , i <- [0 .. pred n], not $ testBit bmp i
      , let cs   = scanl (+)   cost $ zipWith (\b a -> abs $ b - a) bs1 (drop i as)
      , let bmps = scanl setBit bmp $ takeWhile (not . testBit bmp) [i .. pred n]
      , p <- zip bmp cs
      ]
```

しかしTLE,MLEする。

空間計算量は、配列の要素数で$O(2^N)$である。
時間計算量は、`bmp`のループは$2^N$回、`distrib`の中で開始点`i`を$1$から$N$まで回し、終了点も$N$まで探しに行く可能性があるので、$O(N^2 2^N)$で、これは[フェネックさんのいうところの](https://twitter.com/kyopro_friends/status/1723336558978511287)『時間O(N^2 2^N)の想定**TLE**解法』なのか？

でも、[公式解説](https://atcoder.jp/contests/abc328/editorial/7644)の日本語の説明はやっぱり目が滑るけれど、実装例（これもC++のなんだか一般的でなさそうな書き方がいっぱいで目が滑るが）の`for`は3重になっていて、やっていることは同じに見える。けれど、$O(N 2^N)$だと言っている。そしてまたその説明がよくわからん…

配る側でなく、配られる側から考えて、いくつの値が配られるかを考える。
ある状態$S$に集まってくる値は、長さ$1$から長さ$N$まで、何らかの長さ$k$の区間について、$S$で$1$になっている連続区間。
全ての$S$について、ある長さ$k$の1になっている区間を選び出す場合の数は、「連続する$k$個の1」という棒一本と、任意の「0のタイル」「1のタイル」を並べて長さ$N$の数字列を作る方法に等しいので、$2^{N-k}$通り。
ひとつの、長さ$k$の$A$の区間が$B$に対応する位置は$N - k + 1$箇所ある。
よって、配る回数=配られる回数の合計は、$\sum_{k=1}^N 2^{N-k} (N - k + 1) = 2^N (N-1)$ なので、$O(N 2^N)$ なるほど。

<details><summary>計算</summary>

$\sum_{k=1}^N 2^{N-k} (N - k + 1) = \sum_{j=0}^{N-1} 2^j (j+1) = T[N]$ とおく。
$T[N] = 2^0 * 1 + 2^1 * 2 + 2^2 * 3 + ... + 2^{N-1} N$
$2T[N] =          2^1 * 1 + 2^2 * 2 + ... + 2^{N-1} (N-1) + 2^N N$
上から下を引くと左辺は $T[N] - 2T[N] = - T[N]$
右辺は $2^0 + 2^1 + 2^2 + ... + 2^{N-1} - 2^N N = 2^N - 1 - 2^N N = 2^N (1 - N) - 1$
</details>

## 結果

`distrib`の内包表記を諦めて、全てを命令的なループで計算する。

```haskell
abc328g :: Int -> Int -> [Int] -> [Int] -> IO Int
abc328g n c as bs =
  do
    costs <- MUV.replicate (succ all1) maxBound
    MUV.write costs 0 (negate c)
    forM_ [0..pred all1] (\bmp -> do
      let bs1 = drop (popCount bmp) bs
      cost <- (c +) <$> MUV.read costs bmp
      forM_ (filter (not . testBit bmp) [0 .. pred n]) (\i -> do
        foldM_ (\(bmp1, cost1) (i1, ai1, bj1) -> do
          let bmp2 = setBit bmp1 i1
          let cost2 = cost1 + abs (bj1 - ai1)
          MUV.modify costs (min cost2) bmp2
          return (bmp2, cost2)
          ) (bmp, cost) (zip3 (takeWhile (not . testBit bmp) [i .. pred n]) (drop i as) bs1)
        )
      )
    MUV.read costs all1
  where
    all1 = 2 ^ n - 1 :: Int
```

これなら間に合った。[1752ms 44MB](https://atcoder.jp/contests/abc328/submissions/47563873)
なんかくやしい。
（$A$を配列でアクセスするともうすこし速くできる。[1152ms 44MB](https://atcoder.jp/contests/abc328/submissions/47544459)）

<!--
## 追加の解説

[空間計算量の改善 by hirayuu_At](https://atcoder.jp/contests/abc328/editorial/7678)の$DP1[S][last]$は、最初のアプローチの配列と同じ形をしている。これを単独で持たず、$DP2[S]$も併用し、かつ前者は工夫した持ち方をする必要がある、ということらしい。
が、こちらはまたこちらで、実装例を追う気になれないので、ここで終わり。
-->

----

# 追記(11/16)

## E 問題について

Shuffleは興味深いけれど、$_nC_k$ を計算する、mutable arrayがあるときに有用な、コンパクトで風変わりな方法という以上の意味を見いだせなかった。

たかだか8要素なので、Union-Findでなく、隣接行列をビット表現してN-1乗まですることで全域木か判定してみたが、Union-Findの方が速い感じだった。

## F 問題について

### 重み付きUnion-Find

「重み」というより「高低差」と呼ぶ方がイメージがつきやすいかも。

Union-Findは、それぞれのノードについて、親のノードの番号を表にする。
それに加えて、親のノードとの高低差も、表で記録しておく。

```haskell
type UnionFind = (IM.IntMap Int, IM.IntMap Int) -- 親頂点、親との高低差

newUF :: UnionFind
newUF = (IM.empty, IM.empty)
```

分割の代表元を取り出す計算で、親のノードを順に辿るとき、高低差も同時に累積することで、代表元との高低差も取り出せる。

```haskell
-- getRoot ufht i = (a,h) : iの代表元aと、aからiの高さh
getRoot :: UnionFind -> Int -> (Int, Int)
getRoot (uf, ht) i = loop i 0
  where
    loop j h =
      case IM.lookup j uf of
        Just k -> loop k (h + ht IM.! j)
        _      -> (j, h)
```

統合を行う計算では、まず、与えられた二つのノードについて、その代表元とそこからの高低差を得る。
両者の代表元が同一なら、高低差の基準が同じなので、そのまま差をとると、与えられた高低差が正しいか判定できる。
両者の代表元が一致しないとき、指定された高低差になるように、代表元どうしの高低差を決めて接続する。

```haskell
-- uniteUF uf i j d : iとjを連結にし、iからjの高さをdにする 既に連結のときは、高さが一致するか判定
uniteUF :: UnionFind -> Int -> Int -> Int -> (UnionFind, Bool)
uniteUF ufht@(uf, ht) i j d
  | a == b    = (ufht, hbj - hai == d) -- iとjは連結で、iからjの高さはhbj - hai
  | otherwise = ((uf1, ht1), True)     -- aを代表元にし、bを接続する aからbの高さを設定する
  where
    (a,hai) = getRoot ufht i -- hai : aからiの高さ
    (b,hbj) = getRoot ufht j -- hbj : bからjの高さhbj
    uf1 = IM.insert b a uf
    ht1 = IM.insert b (hai + d - hbj) ht
```

（裏で）経路圧縮を行い親を付け替えるとき、同時に高低差も正しい値を保つように調整する。
これを使うと、この問題の本体部分は極めて簡潔にできる。

```haskell
abc328f :: Int -> Int -> [[Int]] -> [Int]
abc328f n _q abds = loop newUF 1 abds
  where
    loop _ _ [] = []
    loop uf i ((a:b:d:_):abds) =
      case uniteUF uf b a d of
        (uf1, True ) -> i : loop uf1 (succ i) abds -- 高低差が矛盾しないか、新規の連結のときは番号を出力
        (uf1, False) ->     loop uf1 (succ i) abds -- 高低差に矛盾がある番号は出力しない
```

## G問題について

[空間計算量の改善 by hirayuu_At](https://atcoder.jp/contests/abc328/editorial/7678)のMLE版のコードを読み解こうとしたのだが、
$DP1[S][last]$ を設定している 31 行め

```cpp
				dp[i][j+1]=min(dp[i^(1<<j)][j],mini[i^(1<<j)]+C)+abs(A[b-1]-B[j]);
```

と、説明

> $DP1[S][last]=S$ に含まれる要素までをすでに確定させていて、最後の要素を $last$ で使うときのコストの最小値

から、添え字の意味がどうにも理解できなくて、降参。
（last = j について考えているときに、j+1 について考える、しかもj+1がS = iに含まれるかどうかを無視して何を計算している？）

ここさえ超えれば、

- 通常のビットDPでは、ビット表現を整数の昇順で計算していけば、注目している状態の部分集合は全て計算済みという性質だけを使うため、それら全体をメモリに保持し、整数のカウントアップの順に計算していく。
- ここでの手法は、要素が1だけ少ない＝ビットが1つだけ落ちた要素だけを参照するので、popCountの数ごとにステップを分けて、不要になったステップの情報は捨てる。

というアイデアでメモリの節約をしているのだろうと推測したのだけれど。