ABC375をHaskellで

ご無沙汰になってしまいました。

# A - Seats

[問題 ABC375A](https://atcoder.jp/contests/abc375/tasks/abc375_a)

シグネチャを決める。

```haskell
import Data.List

abc375a :: Int    -- N
        -> String -- S
        -> Int    -- 答え
abc375a _n s = length $ filter id $ zipWith3 p s (tail s) (drop 2 s)
  where
    p '#' '.' '#' = True
    p _   _   _   = False
```

文字列は長いが `ByteString` でなくても時間の問題はなかった。

# B - Traveling Takahashi Problem

[問題 ABC375B](https://atcoder.jp/contests/abc375/tasks/abc375_b)

シグネチャを決める。

```haskell
abc375b :: Int     -- N
        -> [[Int]] -- Xi, Yi
        -> Double  -- 答え
abc375b _n xys = sum $ zipWith d (o : xys) (xys ++ [o])
  where
    o = [0,0]
    d xy zw = sqrt $ fromIntegral $ sum $ map (^ 2) $ zipWith (-) xy zw
```

# C - Spiral Rotation

[問題 ABC375C](https://atcoder.jp/contests/abc375/tasks/abc375_c)

シグネチャを決める。

```haskell
abc375c :: Int      -- N
        -> [String] -- Aij
        -> [String] -- 答え
```

ややこしいが、外周から数えて、$k$ 層めは $k \bmod 4$ 回、時計回りに90度回転をする。
なので、1～3ステップ回転させた画像を作り、それぞれのピクセルの内容を、層によって妥当な画像から取り出す。

## 結果

```haskell
import Data.Array.Unboxed
import Data.List.Split
import Data.Bool

abc375c :: Int -> [String] -> [String]
abc375c n ss = chunksOf n $ map (bool '.' '#') $ elems t
  where
    bnds = ((1,1),(n,n))
    t0 = listArray bnds $ map ('#' ==) $ concat ss :: UArray (Int,Int) Bool
    t1 = rot t0
    t2 = rot t1
    t3 = rot t2
    rot t = array bnds [((y, succ n - x), a) | ((x,y), a) <- assocs t]
    t = array bnds
        [ (xy, tsi ! xy)
        | i <- [1 .. div n 2]    -- 層の番号
        , let tsi = [t0, t1, t2, t3] !! mod i 4
        , j <- [i .. succ n - i] -- その層の座標範囲
        , xy <- [(i, j), (j, i), (succ n - i, j), (j, succ n - i)]
        ] :: UArray (Int,Int) Bool
```

# D - ABA

[問題 ABC375D](https://atcoder.jp/contests/abc375/tasks/abc375_d)

シグネチャを決める。

```haskell
abc375d :: String -- S
        -> Int    -- 答え
```

中央の要素 $j$ に注目して考える。
$j$ より左のそれぞれの種類ごとの文字の個数を $L_j[c]$ 、より右の文字の個数を $R_j[c]$ とすると、$j$ を中心として回文になっている位置の組 $(i,j,k)$ の個数は
$$ \sum_{'A' \leq c \leq 'Z'} L_j[c] \cdot R_j[c]$$
注目点を一つ進めると、$L_{j+1}[c]$ には $S_j$ が1増える。$R_{j+1}[c]$ には $S_{j+1}$ が1減る。
$L_1[c] = 0$, $R_1[c]$ は $S$ 全体から $S_1$ を引いたカウントになる。

## 結果

```haskell
import Data.Array.Unboxed

abc375d :: String -> Int
abc375d s = ans
  where
    ci0 = listArray ('A','Z') $ repeat 0 :: UArray Char Int
    ck0 = accum (+) ci0 [(c,1) | c <- s]
    (_,_,ans) = foldl' step (ci0, ck0, 0) s
    step (ci, ck, acc) sj = (ci1, ck1, acc + cnt)
      where
        cnt = sum $ zipWith (*) (elems ci) (elems ck1)
        ci1 = accum (+) ci [(sj,  1)]
        ck1 = accum (+) ck [(sj, -1)]
```

# E - 3 Team Division

[問題 ABC375E](https://atcoder.jp/contests/abc375/tasks/abc375_e)

シグネチャを決める。

```haskell
abc375e :: Int     -- N
        -> [[Int]] -- Ai, Bi
        -> Int     -- 答え
```

[フレンズさんを頼る](https://x.com/kyopro_friends/status/1845098688286245325)
> アライグマ「E問題はDPなのだ！
DP[i][x][y]=i番目のメンバーまでで、チーム1の強さをx、チーム2の強さをyにするための、変更する人数
とすればいいのだ！」

$cnt[i][x][y]$ を、$i$ 番目のメンバーまでで、チーム1の強さを $x$ 、チーム2の強さを $y$ にするための変更する最小人数 とする。ただし、不可能な場合は無限大をとる。

3つめのチームは合計の残りなので、追跡する必要はない。
ただし $i$ 人目までの累積和は持っておく必要がある。$BSum_i = \sum_{j=1}^i B_j$ とする。

$cnt[i][x][y]$ の $x,y$ の範囲は $0 \leq x, 0 \leq y, x + y \leq BSum_i$ 
$i = 0$ に対しては $cnt[0][0][0] = 0$ だけ。

以降、何らかの $p,q$ について $C = cnt[i-1][p][q]$ としたとき、

- $A_i = 1$ のとき、人 $i$ を
  - そのまま1に入れると $cnt[i][p+B_i][q] = C$ になる
  - 2に移籍させると $cnt[i][p][q+B_i] = C + 1$ になる
  - 3に移籍させると $cnt[i][p][q] = C + 1$ になる

それぞれの $p,q$ について上の可能性の中の最小値をとると、実際の $cnt[i][x][y]$ が得られる。
同様に、

- $A_i = 2$ のとき、人 $i$ を
  - そのまま2に入れると $cnt[i][p][q+Bi] = C$
  - 1に移籍させると $cnt[i][p+Bi][q] = C + 1$
  - 3に移籍させると $cnt[i][p][q] = C + 1$
- $A_i = 3$ のとき、人 $i$ を
  - そのまま3に入れると $cnt[i][p][q] = C$
  - 1に移籍させると $cnt[i][p+Bi][q] = C + 1$
  - 2に移籍させると $cnt[i][p][q+Bi] = C + 1$

最後に $3m = BSum_N$ について $cnt[N][m][m]$ が答え。

## 結果

```haskell
import Data.Array.Unboxed

abc375e :: Int -> [[Int]] -> Int
abc375e n abs
  | bmod3 /= 0    = -1
  | ans >= tooBig = -1
  | otherwise     = ans
  where
    (bsum3, bmod3) = divMod (sum $ map (!! 1) abs) 3
    tooBig = div maxBound 2                           -- n+1 で十分
    initial = listArray ((0,0),(0,0)) [0] :: UArray (Int,Int) Int
    (final, _bsum) = foldl' step (initial, 0) abs
    ans = final ! (bsum3, bsum3)
    step (cnt, bacc) (a:b:_) = (cnt1, bacc1)
      where
        bacc1 = bacc + b
        cnt1 = accumArray min tooBig ((0, 0),(bacc1, bacc1)) $
               concat $ f a b $ assocs cnt
    f 1 b pqcs = [ [((p+b,q), c), ((p,q+b), succ c), ((p,q), succ c)]
                 | ((p,q),c) <- pqcs]
    f 2 b pqcs = [ [((p+b,q), succ c), ((p,q+b), c), ((p,q), succ c)]
                 | ((p,q),c) <- pqcs]
    f 3 b pqcs = [ [pqc, ((p+b,q), succ c), ((p,q+b), succ c)]
                 | pqc@((p,q),c) <- pqcs]
```

公式解説に、$x,y$ の範囲は $m$ まででよい、とあるけど、巨大な $B_i$ で越えて、再度戻ってくるようなパターンもありそうな気がするのだけど。

# F - Road Blocked

[問題 ABC375F](https://atcoder.jp/contests/abc375/tasks/abc375_f)

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

```haskell
abc375f :: [Int]   -- N,M,Q
        -> [[Int]] -- Ai, Bi, Ci
        -> [[Int]] -- query_i
        -> [Int]   -- 答え
```

方針に自信が持てなかったので[フレンズさんを頼る](https://x.com/kyopro_friends/status/1845099230131781977)
> アライグマ「F問題は逆から考えるのだ！　最初にワーシャルフロイドで最短距離を求めておけば、辺が追加されたときにはその辺を通るか通らないかだけを考えればよくて、最短距離の計算し直しはO(N^2)でできるのだ！」

だいたいあってた。

まず、クエリ1で通行止めになる道を全て除いた、最後まで通行できる道路だけでの、各都市間の距離を求める。これはワーシャルフロイド法の出番。得られた都市 $i$ から都市 $j$ までの距離を $d[i][j]$ とする。

時計を逆回しして、後ろからクエリに対応する。
クエリ2に対して、距離表をひけば答え $d[x][y]$ が得られる。
クエリ1に対して、道路 $i$ が都市 $A_i, B_i$ 間を距離 $C_i$ で結ぶようになったので、全ての都市対 $1 \leq p, q \leq N$ について、新たな距離表 $d'$ が
$$d'[p][q] = \min(d[p][q], d[p][A_i] + C_i + d[B_i][q], d[p][B_i] + C_i + d[A_i][q])$$ として $O(N^2)$ で得られる。

## 結果

```haskell
import Data.Array.Unboxed

abc375f :: [Int] -> [[Int]] -> [[Int]] -> [Int]
abc375f [n,m,_q] abcs qis = ans
  where
    abca = listArray (1,m) abcs :: Array Int [Int]
-- 最後まで残る道の番号
    rs = accumArray (&&) True (1,m) [(i, False) | 1:i:_ <- qis] :: UArray Int Bool
-- 最後まで残る道だけのグラフ
    dm = warshallFloyd n $ [((i,i),0) | i <- [1 .. n]] ++
         concat [[((a,b),c),((b,a),c)] | (True, a:b:c:_) <- zip (elems rs) abcs]
    (_,ans) = foldr step (dm, []) qis
    step (2:x:y:_) (dm, ans) = (dm, (if dmxy == maxBound then -1 else dmxy) : ans)
      where
        dmxy = dm ! (x,y)
    step (1:i:_) (dm, ans) = (dm1, ans)
      where
        a:b:c:_ = abca ! i
        dm1 = array ((1,1),(n,n)) $ concat
              [ [((i,j),d),((j,i),d)]
              | i <- [1 .. n], j <- [i .. n]
              , let [dij,dia,dib,daj,dbj] = map (dm !) [(i,j),(i,a),(i,b),(a,j),(b,j)]
              , let d = minimum $ dij : [dia + c + dbj | dia < maxBound, dbj < maxBound]
                                     ++ [dib + c + daj | dib < maxBound, daj < maxBound]
              ]

warshallFloyd :: Int                  -- 頂点数
              -> [((Int,Int),Int)]    -- 辺の重み
              -> UArray (Int,Int) Int -- 最小コスト
warshallFloyd n edges = dn
  where
    range = [1..n]
    d0 = accumArray (flip const) maxBound ((1,1),(n,n)) edges
    dn = foldl step d0 range
    step d k = d //
      [((i,j), dikj)
      | i <- range, let dik = d ! (i,k), dik < maxBound
      , j <- range, let dkj = d ! (k,j), dkj < maxBound
      , let dikj = dik + dkj, d ! (i,j) > dikj
      ]
```

# G - Road Blocked 2

[問題 ABC375G](https://atcoder.jp/contests/abc375/tasks/abc375_g)

シグネチャを決める。

```haskell
abc375g :: Int     -- N
        -> Int     -- M
        -> [[Int]] -- Ai, Bi, Ci
        -> [Bool]  -- 答え
```

[頼る。](https://x.com/kyopro_friends/status/1845099710228594765)
> アライグマ「G問題は、最短経路に使う辺だけを残したグラフで橋の列挙をすればいいのだ！　グラフの形が特殊だから、累積和みたいにして解く方法もあるのだ」

まだよくわからなくて、公式解説も読んでなぞる形になった。

頂点 $a$ から $b$ までの距離を $d[a][b]$ とする。

ダイクストラ法で1からNまでの最短距離を求めると、副産物として、最短経路のひとつは得られるが、全てではない。これで $d[1][\*]$ が得られる。
Nから1まで戻ることで最短経路に関わる辺と頂点を調べ上げることもできるが、少々めんどくさい。
もう一度ダイクストラ法で、Nからの距離 $d[N][\*]$ を求めると、辺 $(A_i, B_i, C_i)$ について、
- $d[1][A_i] + C_i + d[N][B_i] = d[1][N]$ ならば、最短経路に $i$ は順方向で使われる
- $d[1][B_i] + C_i + d[N][A_i] = d[1][N]$ ならば、最短経路に $i$ は逆方向で使われる
- どちらでもないとき、最短経路に $i$ は関与しない

とわかる。

最短経路に関与する頂点と辺だけからなるグラフを想像する。1とNを両手で持って引っ張ると、全ての辺がたるむことなく、頂点 $a$ は左から距離 $d[1][a]$ の位置に来る。任意の位置での断面を考えると、辺が1本だったり複数本だったりする。複数本な箇所でそのうち1本を切っても、他の最短経路が残る。辺が1本だけの位置で切ると、1とNは切り離されてしまう。

このような、最短経路に関与していて、かつバックアップの最短経路がないような辺だけがこの問題でYesと返すべきものである。

辺を全て調べ、最短経路に関与しているとき、1に近い方の頂点Aについて、距離 $d[1][A]$ をキー、辺の番号 $i$ を値として IntMap に貼り付ける。遠い方の頂点Bについて、距離 $d[1][B]$ をキー、$-i$ を値として貼り付ける。

このIntMapを距離0から、現在走っている辺の番号集合を累積和で追跡する。頂点から抜けたときに、辺集合がsingletonであれば、その要素が探していた辺の番号である。

## 結果

```haskell
import Data.Array.Unboxed

import qualified Data.Heap as H
import Control.Monad.ST
import Data.Array.ST

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

abc375g :: Int -> Int -> [[Int]] -> [Bool]
abc375g n m abcs = elems ba
  where
    g :: Array Int [(Int,Int)]
    g = accumArray (flip (:)) [] (1,n) $ concat
        [[(a,(b,c)),(b,(a,c))] | a:b:c:_ <- abcs]
    d1, dn :: UArray Int Int
    d1 = runSTUArray $ dijkstra n (g !) 1
    dn = runSTUArray $ dijkstra n (g !) n
    d1n = d1 ! n -- 最短距離
-- 辺が最短経路に関わるか判定し、両端の頂点の1からの距離に、自分の番号が出辺か入辺か区別して登録する
    im = IM.fromListWith (++) $ concat
         [ if fw then [(d1a, [i]), (d1b, [-i])] else [(d1b, [i]), (d1a, [-i])]
         | (i, a:b:c:_) <- zip [1..] abcs
         , let dnc = d1n - c, let d1a = d1 ! a, let d1b = d1 ! b
         , let fw = d1a + dn ! b == dnc -- 順方向
         , let bk = d1b + dn ! a == dnc -- 逆方向
         , fw || bk
         ]
-- 1からの距離0から始めて、次の頂点を出発したときの頂点数が1なら橋なので答えとして出力
    (_imN, mes) = mapAccumL step IS.empty $ IM.elems im
    step is es
      | IS.null is2 = (is2, Nothing)
      | u == v      = (is2, Just u )
      | otherwise   = (is2, Nothing)
      where
        is1 = IS.difference is  (IS.fromList $ map negate $ filter (0 >) es)
        is2 = IS.union      is1 (IS.fromList $ filter (0 <) es)
        u = IS.findMin is2
        v = IS.findMax is2

-- 答えとして、挙げられた辺番号だけTrueなリストを作る
    ba :: UArray Int Bool
    ba = accumArray (||) False (1,m) [(e,True) | Just e <- mes]

-- @gotoki_no_joe
dijkstra :: Int                  -- 頂点数N (1～N)
         -> (Int -> [(Int,Int)]) -- 隣接頂点とその辺の重み、グラフの情報
         -> Int                  -- 開始点
         -> ST s (STUArray s Int Int)
dijkstra n edges start =
  do
    dist <- newArray (1,n) maxBound
    writeArray dist start 0
    loop dist $ H.singleton (H.Entry 0 start)
  where
    loop dist queue | H.null queue = return dist
    loop dist queue = do
      let Just (H.Entry cost u, queue1) = H.uncons queue
      du <- readArray dist u
      if du < cost then loop dist queue1 else do
        queue2 <- foldM (\q (v, we) -> do
          let d1 = du + we
          dv <- readArray dist v
          if d1 >= dv then return q else do
            writeArray dist v d1
            return $ H.insert (H.Entry d1 v) q
          ) queue1 (edges u)
        loop dist queue2
```

`Data.IntSet.size` は $O(N)$ なので、要素が1であることを判別するのに、最小値と最大値が等しい、という計算を書いてみた。個数を別に追跡するより楽ちんだったので。

`Data.IntMap.alter` はあるのに `Data.IntSet.alter` がない（`alterF` はある）のはなぜだろう。
Identityモナドで包んで使ってみたら[上より速くなった。](https://atcoder.jp/contests/abc375/submissions/58906156)

# おわりに

immutable array が活躍する回で楽しかった。