ABC320 A～F をHaskellで

A - Leyland Number

問題 ABC320A

結果

abc320a :: Int -- A
        -> Int -- B
        -> Int -- 答え
abc320a a b = a ^ b + b ^ a

B - Longest Palindrome

問題 ABC320B

シグネチャを決める。

abc320b :: String  -- S
        -> Int     -- 答え

外から調べる

$2 \leq |S| \leq 100$ と大したサイズではないので、長さ $|S|$ から 1 まで順に、全ての開始位置から調べて、その長さの回文があるかを調べ、最初に該当した長さを答えとする、でできる。
ランダムアクセスをしてしゃれた作りにしてみる。

import Data.Array

abc320b :: String -> Int
abc320b s = head [k | k <- [n, pred n..1], (i,j) <- zip [1..] [k..n], check i j]
  where
    n = length s
    a = listArray (1,l) s
    check i j
      | i >= j         = True
      | a ! i /= a ! j = False
      | otherwise      = check (succ i) (pred j)

中から調べる

逆に、あらゆる位置から開始して、そこを中心とした回文の長さを集めて、最大値をとる方策もある。
奇数長と偶数長の両方を調べる必要がある。

abc320b s = maximum [pallen i j | i <- [1..n], j <- [i, succ i]]
  where
    n = length s
    a = listArray (1,n) s
    pallen i j
      | 1 <= i, j <= n, a ! i == a ! j = pallen (pred i) (succ j)
      | otherwise = pred j - i

C - Slot Strategy 2 (Easy)

問題 ABC320C

シグネチャを決める。

abc320c :: Int     -- M
        -> String  -- S1
        -> String  -- S2
        -> String  -- S3
        -> Int     -- 答え

同時には１つのボタンしか押せないので、狙っている数字が同時に複数のリールに見えているときにどれを選ぶかで結果が変わるということ。

リールは3つなので、「うまい選び方を見つける」代わりに、$3! = 6$ 通りの順序で全てやってみて最短を見つける、という力任せの解法をとる。どの数字を狙うかも総当たりする。

結果

abc320c :: Int -> String -> String -> String -> Int
abc320c m s1 s2 s3
  | null cands = -1
  | otherwise  = minimum cands
  where
    cands =
      [ findmin pts
      | d <- ['0'..'9']             -- 数字を全て試す
      , let tss = map (elemIndices d) [s1,s2,s3] -- 数字の出現位置
      , not (any null tss)          -- 数字dが出現しないリールがあると不可能
      , pts <- permutations tss     -- リールの押す順序を総当たり
      ]
    findmin [ts1,ts2,ts3] = t3
      where
        t1 = head ts1  -- リール1を止める
        t2 = head $ dropWhile (t1 >=) $ cycleM ts2 -- t1以降の時刻でリール2を止める
        t3 = head $ dropWhile (t2 >=) $ cycleM ts3 -- t2以降の時刻でリール3を止める
    cycleM ts = ts ++ cycleM (map (m +) ts) -- 出現位置をループさせる

「楽な解法」

楽な解法 by evima いわく、「揃えられるならリール3周以内に達成できるからその総当たりできる」の方針で、なるべく「楽をする」方向で書いてみた。

abc320c :: Int -> String -> String -> String -> Int
abc320c m s1 s2 s3
  | null cands = -1
  | otherwise  = minimum cands
  where
    cands =
      [ maximum [i,j,k]
      | (i, c1) <- zip [0..] $ s1 ++ s1 ++ s1
      , (j, c2) <- zip [0..] $ s2 ++ s2 ++ s2, i /= j,         c1 == c2
      , (k, c3) <- zip [0..] $ s3 ++ s3 ++ s3, i /= k, j /= k, c1 == c3
      ]

D - Relative Position

問題 ABC320D

シグネチャを決める。A,B,X,Yは手抜きする。
結果は、長さ2のリストのとき座標、空リストのとき undecidable を表す。

abc320d :: Int      -- N
        -> [[Int]]  -- A,B,X,Y
        -> [[Int]]  -- 答え

「人 $A_i$ から見て、人 $B_i$ は、$x$ 軸正方向に $X_i$、$y$ 軸正方向に $Y_i$離れた位置にいる」
という情報を、人を頂点とし、相対関係を有向辺とするグラフにする。逆向きの辺も張る。
このグラフを、原点にいる人1を起点として、ここから到達可能な人物は座標が定められる。

結果

簡単に、深さ優先探索して、結果は IntMap に追記する。

import Data.Array
import qualified Data.IntMap as IM

abc320d :: Int -> Int -> [[Int]] -> [[Int]]
abc320d n m abxys = IM.elems xyM
  where
    g = accumArray (flip (:)) [] (1,n)         -- いつものグラフ
        [p | a:b:x:y:_ <- abxys, p <- [(a,(b,x,y)), (b,(a,-x,-y))]]
    xy0 = IM.fromAscList [(k,[]) | k <- [1..n]]   -- 全ての人がundecidableと書かれた表
    xyM = loop xy0 [(1,0,0)]       -- 人1を原点に配置するところから、深さ優先探索で表を埋める
    loop xyM [] = xyM
    loop xyM ((k,x,y):kxys)
      | not $ null $ xyM IM.! k = loop xyM kxys -- 人kは確定済み
      | otherwise = loop (IM.insert k [x,y] xyM) (kxys1 ++ kxys)
      where
        kxys1 = [(b,x+dx,y+dy) | (b,dx,dy) <- g ! k, null (xyM IM.! b)]

E - Somen Nagashi

問題 ABC320E

シグネチャを決める。$T_i,W_i,S_i$は手抜きする。

abc320e :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- T,W,S
        -> [Int]    -- 答え

各時刻で、そうめんを取るためにスタンバっている人の集合を管理する。IntSet は最小要素を定数時間で取り出せる。

この世界では「そうめんが流れてくる」と「人が列に復帰する」の２つのイベントが発生する。

data Event = ComeBack Int  -- ComeBack i = 人iが復帰する
           | Somen Int Int -- Somen w t = 量wのそうめんが流れる。これを得た人は時刻 t に復帰する

入力データにおいて、そうめんの時刻は整列されているが、人が復帰する時刻は自由なので、時刻をキーとする優先度付きキューでイベントを管理して、シミュレーションを実行する。
人の復帰とそうめんが同時刻のときは間に合うという仕様なので、復帰が先に処理されるようにする必要がある。人の復帰時刻 $t$ をキューの優先度 $2t$、そうめんの時刻 $t$ をキューの優先度 $2t + 1$と最下位ビットに埋め込んでしまうのが楽か。

結果

import Data.Array
import qualified Data.IntSet as IS
import qualified Data.Heap as H

data Event = ComeBack Int  -- ComeBack i = 人iが復帰する
           | Somen Int Int -- Somen w t = 量wのそうめんが流れる。これを得た人は時刻 t に復帰

abc320e :: Int -> Int -> [[Int]] -> [Int]
abc320e n _m twss = elems acc
  where
    q0 = H.fromList [H.Entry (t + succ t) (Somen w (t+s)) | t:w:s:_ <- twss]
    s0 = IS.fromList [1..n]
    acc = accumArray (+) 0 (1,n) $ loop q0 s0

loop :: H.Heap (H.Entry Int Event) -> IS.IntSet -> [(Int, Int)]
loop q s =
  case H.uncons q of
    Nothing -> []
    Just (H.Entry _ (ComeBack i), q1) -> loop q1 (IS.insert i s)
    Just (H.Entry _ (Somen w t ), q1) ->
      case IS.minView s of
        Nothing -> loop q1 s
        Just (i, s1) -> (i, w) : loop (H.insert (H.Entry (t + t) (ComeBack i)) q1) s1

公式解説から

セグメント木を用いる解法 by sounansyaによると、

それぞれの参加者が隊列に復帰する時刻をセグメント木に入れる。演算は $\min$ にする。
そうめんを得た人の復帰時刻でセグメント木を更新する。
そうめんを得る人は、復帰時刻がそうめんが流れてくる時刻以下の、背番号最小の人。セグメント木に入れてあるので、任意の背番号区間について、該当者がいるかどうかは高速に判定できる。そこで二分探索することで、背番号最小の該当者を特定できる。

素直にセグメント木を使うと、二分探索の手数が$O(\log N)$で、区間の最小背番号を求めるのに毎回$O(\log N)$かかるため、結局$O(\log^2 N)$かかることになる。
自由に区間を指定して外から二分探索するのではなく、セグメント木自身が二分木であることを利用して、この問題の答えを直接求めるような計算が atcoder::segtree.max_right() ということだろうか。

F - Fuel Round Trip

問題 ABC320F

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

abc320f :: Int      -- N
        -> Int      -- H
        -> [Int]    -- Xi
        -> [[Int]]  -- Pi,Fi
        -> Int      -- 答え

わからなくて公式解説を見た。

位置$i$ ($0 \leq i \leq N$) を座標 $X_i$ と呼ぶことにする。ただし $X_0 = 0$ とする。
位置$i-1$から位置$i$までの距離$X_i-X_{i-1} = D_i$ と定義する。

各ガソリンスタンドは、往路で使う、復路で使う、使わずじまい、の3通りの選択肢がある。
その組み合わせで、戻ってこれる最小コストを探すことが目的である。

$cost[i,j,k]$ を、位置$i$で往路では$j$リットル、復路では$k$リットルの燃料を持っているときに、位置$1～i$で支払う金額の合計の最小値、とするDPを行う。
ただし、位置$i$のガソリンスタンドを使用する場合、

• $j$ : 往路においては使用した後の値
• $k$ : 復路においては使用する前の値

とする。これは、$i = 0$からDPを進めていったとき、注目している位置のスタンドについてどうするかをそのときに決定できるようにするためである。

位置0での$cost$は、燃料は満タンにして出発するため、$cost[0,H,j] = 0, cost[0,i \neq H, j] = \infty$ とするのが正しいが、オール0でも結果は変わらない。

位置$i-1$での$cost[i-1,j,k]$と、次の位置$i$ までの距離$D_i$、位置$i$ のスタンドの価格 $P_{i+1}$ 補給量 $F_{i+1}$ から、位置$i$での$cost[i,j,k]$ を求める。

まず、位置$i$のスタンドは使わない場合を考える。これを$c0[j,k]$とする。
位置$i-1$からの距離$D_i$で消費する分を考えて、

c0[j,k] = \left \{
\begin{array}{ll}
cost[i-1,j+D_i, k-D_i] & j+D_i \leq H, k - D_i \geq 0 のとき \\
\infty & それ以外のとき
\end{array} \right .

となる。
往路でスタンドを使う場合を$c1[j,k]$とする。
スタンドを使う前に位置$i$には到達できている必要があるので、$c0$の$j$を$F_i$だけずらした位置の値を参照すればよい。ただし、$j=H$については、燃料が入りきらなかった可能性もある。

c1[j,k] = \left \{
\begin{array}{ll}
P_i + c0[j-Fi,k] & j - F_i \geq 0, j < H のとき \\
P_i + \min c0[x, k] & (H-F_i \leq x \leq H), j = H のとき \\
\infty & それ以外のとき
\end{array} \right .

復路でスタンドを使う場合を$c2[j,k]$とする。
$k$からさらに$F_i$だけ補給してから出発するが、入りきらない場合もある。

c2[j,k] =
P_i + c0[j, \min(H, k + Fi)]

これらをpointwiseに最小値をとればよい。

cost[i,j,k] = \min(c0[j,k], c1[j,k], c2[j,k])

最後の位置 $N$ についてはスタンドがないので、$P_N = \infty, F_N = 0$ とでも置けばよい。

位置$N$に到達したときに燃料が$j$あったら、そのまま折り返すので、復路の燃料も$j$である。
ここから、問題の答えは $\min cost[N,j,j] \ (0 \leq j \leq H)$ となる。

結果

3次元配列に全ての状態を記録する必要はなく、直前の状態だけから次の状態は求められるのでそうする。
$c0$は配列で作り、$c1,c2$はそれに accum min する要素のみ生成している。

import Data.Array

tooBig = 300 * 10^5 :: Int

abc320f :: Int -> Int -> [Int] -> [[Int]] -> Int
abc320f n h xs pfs
  | ans < tooBig = ans
  | otherwise    = -1
  where
    ds = zipWith (-) xs (0:xs)
    bnds = ((0,0),(h,h))
    cost0 = accumArray (flip const) tooBig bnds [((h,j), 0) | j <- [0..h]]
    costN = foldl' step cost0 $ zip ds (pfs ++ [[0,0]])
    ans = minimum [costN ! (j,j) | j <- [0..h]]
    step costi (di, pi:fi:_) = accum min costA (costBs ++ costCs)
      where
        costA = listArray bnds $ map cA $ range bnds
        cA (j,k) | j <= h - di, k >= di = costi ! (j + di, k - di)
                 | otherwise = tooBig
        costBs =
          [ ((j,k), pi + minimum [costA ! (j1,k) | j1 <- [h - fi .. pred h]])
          | let j = h, h - fi < h, k <- [di..h]
          ] ++
          [ ((j,k), pi + costA ! (j - fi, k))
          | j <- [fi..pred h], k <- [0..h]
          ]
        costCs =
          [ ((j,k), pi + costA ! (j, min h (k + fi)))
          | j <- [0..h], k <- [0..h]
          ]