ABC369をHaskellで

今回はわかりやすかったです。
時間内に解ける訳ではありませんが。

# A - 369

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_a)

シグネチャを決める。

```haskell
abc369a :: Int -- A
        -> Int -- B
        -> Int -- 答え
```

$A < B$ を仮定して、$x_1 < A < B, A < x_2 < B, A < B < x_3$ と3通りの位置に $x$ が当てはまる。
（$B < A$ のときも同じ）
$x_1$ は $B - A = A - x_1, x_1 = 2A - B$ と必ず決められる。
$x_3$ も $x_3 - B = B - A, x_3 = 2B - A$ となる。
$x_2$ は $B - x_2 = x_2 - A, 2x_2 = A + B$ これは $A+B$ が奇数だと $x_2$ が整数にならないので、条件付きとなる。

仮定に反する $A = B$ のとき、$x_1 = x_2 = x_3$ な1通りしかない。

という考察と対応するテストケースが全て例示されている。やさしいね。

## 結果

```haskell
abc369a a b
  | a == b      = 1
  | odd (a + b) = 2
  | otherwise   = 3
```

# B - Piano 3

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_b)

シグネチャを決める。

```haskell
abc369b :: Int           -- N
        -> [(Int, Char)] -- Ai, Si
        -> Int           -- 答え
```

$S_i$ が `L` か `R` かで $A_i$ を二つに分けて、前後の差の合計をとればよい。
初期位置は、それぞれの手が最初に押すキーの上ということにしておけば、コストを最も節約できる。

## 結果

```haskell
abc369b _n ass = f (getas 'L') + f (getas 'R')
  where
    getas c = [a | (a,s) <- ass, s == c]
    f as = sum $ zipWith d as $ tail as
    d x y = abs $ x - y
```

# C - Count Arithmetic Subarrays

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_c)

シグネチャを決める。

```haskell
abc369c :: Int   -- N
        -> [Int] -- Ai
        -> Int   -- 答え
```

問題文にあるとおり、長さ1の数列は等差数列である。これは別に数える。
また、長さ2の数列も等差数列である。2項あって初めて公差が定まるが、その公差で来るべき次の項がないだけなので。
つまり、項の間の差分 $D_i = A_{i+1} - A_i$ をとり、これが等しい値が $k$ 項連続するとき、$A_i$ の方では $k+1$ 項の等差数列があるとわかる。
この範囲のより短い区間も等差数列になるので、ここには $\_{k+1}C_2 = k(k+1)/2$ 個の（長さ2以上の）等差数列がある。

## 結果

```haskell
abc369c n as = n + m
  where
    m = sum $ map (k1C2 . length) $ group $ zipWith (-) as $ tail as
    k1C2 k = div (k * succ k) 2
```

# D - Bonus EXP

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_d)

シグネチャを決める。Cと同じになってしまった。

```haskell
abc369d :: Int   -- N
        -> [Int] -- Ai
        -> Int   -- 答え
```

次に倒すときそれが奇数回めか偶数回めか、の2つの場合について、その最大スコアを追跡するDPで解ける。
逃がすとき、経験値+0、場合の変化なし
戦うとき、
奇数回めなら+X、次は偶数回め
偶数回めなら+2X、次は奇数回め

## 結果

```haskell
abc369d _n as = uncurry max $ foldl step (0, minBound) as
  where
    step (o, e) a = (max o $ e + a + a, max e $ o + a)
```

# E - Sightseeing Tour

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_e)

シグネチャを決める。うってかわって入力が多い。
$K_i$ は $B_{i*}$ の長さで必要なら得られるので省く。

```haskell
abc369e :: Int     -- N
        -> Int     -- M
        -> [[Int]] -- Ui,Vi,Ti
        -> Int     -- Q
        -> [[Int]] -- Bij
        -> [Int]   -- 答え
```

$K_i$が控えめなのがヒント。

## アイデア1

最大5つの橋について、一度は通ったかまだか、の状況を見分けられる必要があるので、
島と橋のなすグラフを $2^{K_i}$ 個に複製し、初めて通る橋について、それを通過済な状態を表す複製へと接続し直す。通過済な橋は、普通に同一の複製内に接続する。
このグラフにおいて、0個めのコピーの島1から、$2^{K_i}-1$ 個めのコピーの島$N$までの距離をダイクストラ法で数える。

ダイクストラ法の計算量を $O((E+V)\log V)$ として、複製により頂点や辺が増える量とクエリの回数を考えると、少々心許ない。

## アイデア2

島と橋のなすグラフについて、全ての島 $a$ から島へ $b$ の距離 $d[a][b]$ をワーシャルフロイド法で先に測ってしまう。
$K_i$ 本の橋を渡る順序は $K_i !$ とおりあり、さらに、それぞれの橋をどちら向きに渡るかの場合が $2^{K_i}$ とおりある。島1から最初の橋の始点まで、渡った橋の終点から次の橋の始点まで、最後の橋の終点から$N$まで、は $d$ の距離がある。

アイデア1では、使わない島や橋についての計算を几帳面に実行してしまうが、それを$d$でキャンセルできる。
$d$ で移動する経路に、$B_{i*}$ のいずれかが埋まっていて無駄が出るように思えるが、全ての順序全ての方向を試すので、そのようなものを全く含まない最適な選択も総当たりの中にちゃんと含まれるから大丈夫。

## 結果

```haskell
import Data.List
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed

abc369e :: Int -> Int -> [[Int]] -> Int -> [[Int]] -> [Int]
abc369e n m uvts _q bss = map solve bss
  where
    d = runSTUArray $ warshallFloyd n $ concat
        [[(u,v,t), (v,u,t)] | u:v:t:_ <- uvts] ++
        [(i,i,0) | i <- [1..n]]
    b = listArray (1,m) uvts :: Array Int [Int]
    solve bs = t1 + minimum
      [ sum $ zipWith (curry (d !)) vs us
      | pbs <- permutations bs
      , ds <- replicateM k [True, False]
      , let us =     [b ! i !! if d then 0 else 1 | (i,d) <- zip pbs ds] ++ [n]
      , let vs = 1 : [b ! i !! if d then 1 else 0 | (i,d) <- zip pbs ds]
      ]
      where
        k = length bs
        t1 = sum $ map ((!! 2) . (b !)) bs

-- ワーシャルフロイド法
warshallFloyd :: Int                             -- 頂点数 N
              -> [(Int,Int,Int)]                 -- グラフ（有向、頂点1～Nの始点、終点、距離）
              -> ST s (STUArray s (Int,Int) Int) -- 距離
warshallFloyd n graph =
  do
    d <- newArray ((1,1),(n,n)) maxBound
    forM_ graph (\(i,j,w) -> do
      d0 <- readArray d (i,j)
      writeArray d (i,j) $ min d0 w
      )
    forM_ [1..n] (\k ->
      forM_ [1..n] (\i -> do
        dik <- readArray d (i,k)
        when (dik < maxBound) (
          forM_ [1..n] (\j -> do
            dkj <- readArray d (k,j)
            when (dkj < maxBound) $ do
              let dikj = dik + dkj
              dij <- readArray d (i,j)
              when (dikj < dij) $ writeArray d (i,j) dikj
            )
          )
        )
      )
    return d
```

タイムは 3413ms、制限時間 4sec だったので間に合った。

# F - Gather Coins

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_f)

シグネチャを決める。
出力がいつもより凝っている。

```haskell
abc369f :: Int           -- H
        -> Int           -- W
        -> [[Int]]       -- Ri,Ci
        -> (Int, String) -- 答え
```

$\_nC_k$ を数えるやり方では $H,W$ が大きくて全然間に合わない、とたじろいだが、
よく読んでみると、$N \leq 2 \times 10^5$ ということで、コインの方が制限されている。
マス目でなくてコインに注目して、そのコインを拾っていける経路で最高スコアは、という向きに考てみる。

あるコインに対して、それより左上にあるコインの最大スコアが判れば、それ+1にできる。
左上から行をスキャンする順序で、行ごとに処理していくことを考える。
これまでに調べた位置について、IntMapで情報を保持して、横座標をキー、存在するコインと付随する情報を値にして、次のコインについて、横座標がそれ以下な情報はIntMapで取り出せても、スコア最大のものを見つけるのは線形探索になり間に合わない。

最大値を知りたいので、セグメント木で、最大スコアなコインの座標を取り出す演算を付けてやればいい。
セグメント木の添え字は、コインの横座標。
内容は
- それを選ぶときに得られる最大スコア（それ自身を含む）
- それの位置 (r,c) と、(0,0)からの最大スコアのために経由する位置の全ての逆順のリスト

高速化のために、位置は $(1,1)～(H,W)$を$0～H*W-1$に写した整数で扱う。

## 結果

```haskell
import Control.Monad.ST
import Data.Array.ST

type Elem = (Int,[Int])

abc369f :: Int -> Int -> Int -> [[Int]] -> (Int, String)
abc369f h w n rcs = runST $
  do
-- 空のRQAを作る
    rqa <- makeRQArrayN findmax start w :: ST s (RQArray (STArray s) Elem)
-- コインを順に並べていく
    forM_ cs (\i -> do
      let (_,c) = i2rc i
      (sc, is) <- queryRQArray rqa 0 (succ c)
      updateRQArray rqa c (succ sc, i:is)
      )
-- goalについて最終的な答えを調べる
    (sc, is) <- queryRQArray rqa 0 w
    return (sc, i2cmds $ goal:is)
  where
    rc2i (r:c:_) = w * pred r + pred c -- 座標を辞書順のindexにする
    i2rc i = divMod i w -- indexから座標(-1)に戻す
    cs = sort $ map rc2i rcs -- コイン座標の整列されたリスト
    start = (0, [0]) :: Elem
    goal = rc2i [h, w]
    findmax a b = if fst a >= fst b then a else b
    i2cmds is = concat $ zipWith f ris $ tail ris
      where
        ris = reverse is
        f i j = replicate (s-r) 'D' ++ replicate (d-c) 'R'
          where
            (r,c) = i2rc i
            (s,d) = i2rc j

-- セグメント木
data RQArray ar a = RQA (a->a->a) a Int (ar Int a)

{- makeRQArrayN len op u 要素数を指定して単位元で満たされた区間クエリ配列を作る
要素数len 添え字0始まり
op モノイド演算
e opの単位元
-}
makeRQArrayN :: MArray ar a m => (a->a->a) -> a -> Int -> m (RQArray ar a)
makeRQArrayN op e len = RQA op e w <$> newArray (0, w + pred w) e
  where w = until (len <=) (2 *) 1

{- queryRQArray st a b [a, b)の区間をopした結果を求める -}
queryRQArray :: MArray ar a m => RQArray ar a -> Int -> Int -> m a
queryRQArray (RQA op u w arr) a b = loop arr 0 w 0
  where
    loop arr p w i
      | q <= a || b <= p = return u
      | a <= p && q <= b = readArray arr i
      | otherwise = do
        l <- loop arr  p       w2 (i + succ i)
        r <- loop arr (p + w2) w2 (2 * succ i)
        return (op l r)
      where
        q = p + w
        w2 = div w 2

{- updateRQArray rqa i x 要素iをxに設定する -}
updateRQArray :: MArray ar a m => RQArray ar a -> Int -> a -> m ()
updateRQArray (RQA op _ w arr) j x = updateLoop op arr (j + pred w) x

updateLoop _p arr 0 x = writeArray arr 0 x
updateLoop op arr i x = do
  writeArray arr i x
  y <- readArray arr (if even i then pred i else succ i) -- xor i 1
  updateLoop op arr (div (pred i) 2) (op x y)
```

# G - As far as possible

[問題](https://atcoder.jp/contests/abc369/tasks/abc369_g)

シグネチャを決める。

```haskell
abc369g :: Int      -- N
        -> [[Int]]  -- Ui,Vi,Li
        -> [Int]    -- 答え
```

青木君は、根1から最も遠い葉を指定する。スコアは距離の2倍になる。
続けて、既に踏んだ経路となるべく経路を共有しないような葉を指定する。スコアは新しく踏んだ経路の距離の2倍だけ増える。
ということを、葉の数だけ続ける。葉が尽きたら、中間ノードをどう指定しても以降のスコアは変化しない。

このようなスコアの増分は、それぞれの葉から根までの距離、ただし合流したときは最大のものを残してその他はカウントを止めて降順に持つ、という操作で、1から、または分岐点から葉の距離、を一度に計算できる。

## 結果

距離の降順リストの代わりに、マイナスにした距離の優先度付きキューで保持する。

```haskell
import Data.Array
import qualified Data.Heap as H

abc369g :: Int -> [[Int]] -> [Int]
abc369g n uvls = take n $ ans ++ repeat (last ans)
  where
    g = accumArray (flip (:)) [] (1,n) $ concat
        [[(u,(v,l)),(v,(u,l))] | u:v:l:_ <- uvls]
    recur p v
      | null xs   = H.singleton 0
      | otherwise = foldr1 H.union xs
      where
        xs = [H.adjustMin (subtract l) $ recur v c | (c,l) <- g ! v, c /= p]
    ansH = recur 1 1
    ans = scanl1 (+) $ map ((-2) *) $ unfoldr H.uncons ansH
```