ABC274 A～E+F をHaskellで

入力の規模がいつもより控えめな感じでした。

A - Batting Average

問題 ABC274A

シグネチャを決める。

abc274a :: Int     -- A
        -> Int     -- B
        -> String  -- 答え

四捨五入は前回のBとネタ被り。
小数点第4位で四捨五入する計算を整数で扱うために、$B$ を$10^4$倍してから割る。
さらに10で割った商が小数点第3位から上の数字、余りが第4位の数字。これが5以上なら前者にさらに1を加えて、小数点を加えればよい。
例2の $A=7, B=3$ の場合は次のようになる。

3/7 = 0.42857142857142855…
B * 10^4 = 30000
div ↑ A = 4285
divMod ↑ 10 = 428, 5
四捨五入 → 429
小数点追加 → 0.429

小数点を追加するときは、1の位まで0を追加したりする処理が面倒なので、10の位に1をつけてから show する。
すると、（今回の入力の範囲では存在しないが）四捨五入による繰り上がりで結果が 1.000 になるような場合でもそのまま扱える。

結果

abc274a a b = d : '.' : ds
  where
    (q,r) = divMod (div (b * 10000) a) 10
    (_:d:ds) = show $ 10000 + q + if 5 <= r then 1 else 0

B - Line Sensor

問題 ABC274B

シグネチャを決める。

abc274b :: Int       -- H
        -> Int       -- W
        -> [String]  -- Cij
        -> [Int]     -- 答え

縦の列方向に # がいくつあるかをそれぞれ数えればよい。

結果

リストの transpose と sum を使えば次のようにできる。ただし transpose は行列が大きくなると遅い。

import Data.List

abc274b :: Int -> Int -> [String] -> [Int]
abc274b h w css = map (sum . map c2i) $ transpose css

c2i '#' = 1
c2i '.' = 0

行ごとに c2i した結果を、行ごとに足し合わせていけば同じ結果が得られる。

import Data.List

abc274b :: Int -> Int -> [String] -> [Int]
abc274b h w css = foldr (zipWith (+)) (repeat 0) $ map (map c2i) css

座標を使って accumArray で足し合わせることでも求められる。

import Data.Array.Unboxed

abc274b :: Int -> Int -> [String] -> [Int]
abc274b h w css = elems arr
  where
    arr :: UArray Int Int
    arr = accumArray (+) 0 (1,w)
          [(j,1) | cs <- css, (j,'#') <- zip [1..] cs]

C - Ameba

問題 ABC274C

シグネチャを決める。

abc274c :: Int    -- N
        -> [Int]  -- Ai
        -> [Int]  -- 答え

アメーバは1番から $2N+1$ 番までが登場する。
「何代親を遡るとアメーバ1になるか」の答えは、

• アメーバ1番のそれは0である。
• アメーバ $A_i$ 番が $c$ であるとき、$2i, 2i+1$ 番は $c+1$ である。

これを持つ配列はDP配列と同様の手法で構築できる。

結果

import Data.Array

abc274c :: Int -> [Int] -> [Int]
abc274c n as = elems arr
  where
    arr = array (1, n + n + 1) $ (1, 0) :
      [ p
      | (i, a) <- zip [1..] as
      , let c1 = succ $ arr ! a
      , p <- [(i + i, c1), (i + i + 1, c1)]
      ]

D - Robot Arms 2

問題 ABC274D

シグネチャを決める。

abc274d :: Int    -- N
        -> Int    -- x
        -> Int    -- y
        -> [Int]  -- Ai
        -> Bool   -- 答え

旋回せずに直進することはなく、毎回90度向きを変えるので、つまり $A_i$ の奇数番はX軸方向、偶数番はY軸方向に変化を加える。向きは正方向にも負方向にもできるので、$_i$ は足しても引いてもどちらでもよく（ただし $A_1$ は正方向に限定）、最終結果をどちらもそれぞれ $x, y$ に等しくできるか、と聞いている。

数列 $X_i$ の前 $k$ 個を足すまたは引いて作ることのできる数の集合を $S_k$ とする。
$S_0 = \{0\}$
$S_{k+1} = \{ n + X_k, n - X_k \,|\, n \in S_k \}$
という漸化式をそのまま使うと、$S_{k+1}$ を作るときに全ての要素を計算で作ることになる。
数の集合 $T_k$ とオフセット値 $o$ の対で、$\{ m + o \,|\, m \in T_k \}$ という集合を表現することにすると、$ + X_k$ する代わりにオフセットに $+ X_k$ して、オフセットの増分を考慮して $- 2X_k$ した値を追加することで漸化式と同じ集合を構築できる。

結果

import qualified Data.IntSet as IS

abc274d :: Int -> Int -> Int -> [Int] -> Bool
abc274d n x y (a1:as) = IS.member (x - ox) xs && IS.member (y - oy) ys
  where
    dys = [a | (a, True) <- zip as $ cycle [True, False]]
    dxs = [a | (a, True) <- zip as $ cycle [False, True]]
    (ox,xs) = foldl step (a1, IS.singleton 0) dxs
    (oy,ys) = foldl step ( 0, IS.singleton 0) dys
    step (ofs, ns) x = (ofs + x, IS.union ns $ IS.map (\n -> n - x - x) ns)

E - Booster

問題 ABC274E

シグネチャを決める。

abc274e :: Int          -- N
        -> Int          -- M
        -> [(Int,Int)]  -- Xi,Yi
        -> [(Int,Int)]  -- Pi,Qi
        -> Double       -- 答え

いかにも巡回セールスマン問題で、訪問するべき場所の数も全部で18か所なので集合のビットマップ表現が使える。

番号をそれぞれの地点に割り当てる。ビットマップ表現のビット番号もこれを使う。

番号	地点
$0$ ～ $N-1$	街 $N$ 個
$N$ ～ $N + M - 1$	宝箱 $M$ 個
$N+M$	原点

原点を除く地点について、地点の部分集合を表すビットマップ表現は、空集合が0、全体集合が $2^{N+M} - 1$ である。
訪問済みの地点集合と、その中の最終地点の組に関して、経路の最小コストを与えるDP配列を構築する。
通常は2次元配列で行うが、最終地点は $0$～$N+M-1$ 全体でなくそのうち訪問済みの箇所だけなので、必要な地点番号についてのみ値を持つ IntMap を持つ、ビットマップをキーとする配列でやってみよう。と思ったが、結局、最終地点に対してランダムアクセスする必要がなかったので、IntMap でなく対応付けリストで済む。

地点集合 $bm$ の要素が1つだけのとき、原点からその地点までの距離がコストである。
それ以外のとき、そこに含まれる任意の地点を最終地点 $k$ として、最小コストを求める。
$bm$ から $k$ を除いた地点集合 $bm_1$ について、DP配列から「最終地点ごとの最小コスト」を取り出す。
それぞれの最終地点 $j$ について、そのコストに $j$ から $k$ までのコストを足すと、$j$ を経由して $k$ に至る総コストが得られる。その中の最小値を決めればよい。
なお、「$j$ から $k$ までのコスト」は、$j$ から $k$ までの距離を、$bm_1$ に含まれる宝箱の個数 $t$ により $/ \, 2^t$ した値となる。

「最小値を選ぶ」ことで場合を統合するのがDPのポイントだが、「現時点でのコストでは他に後れを取るが、ブーストにより今後巻き返す」ような場合を無視していると、結果がおかしくなる。しかし、ブースト値は通過した宝箱の個数で決まるので、訪問済みの地点集合が同一ならブースト値は等しく、そういう逆転現象の心配はないとわかる。

最終的に原点に戻る最小コストは、全ての街を訪問した状態から原点へ進んだ場合のコストを同様に求めればよい。ここで、宝箱は取らなくてもよいので、宝箱それぞれを「とった/とっていない」の $2^M$ 通りの地点集合全てについて求めた最小値を選択する必要がある。

結果

地点集合が要素数1の場合を特別対応する代わりに、地点集合が空集合である場合の地点とコストの対応付けリストとして [(n+m, 0.0)] を与えておくことで、一般の場合の計算がそのまま最初の計算と同じになるようにできた。

最後に原点に戻るコストの計算は、DPで次を計算する漸化式と共通するところが多い。というか、特定の最終地点について計算する部分がそのまま使えるので、ここを costfk として括りだした。

import Data.Array

abc274e :: Int -> Int -> Int -> Int -> [(Int,Int)] -> [(Int,Int)] -> Double
abc274e n m xys pqs = ans
  where
    xya = listArray (0,n+m) $ xys ++ pqs ++ [(0,0)]
-- 地点番号から距離
    dist i j = sqrt $ fromIntegral $ (x1 - x2)^2 + (y1 - y2)^2
      where
        (x1,y1) = xya ! i
        (x2,y2) = xya ! j
-- 地点集合からブースト係数
    boost = (0.5 ^) . popCount . (boostMask .&.)
    boostMask = shiftL (2^m - 1) n
-- 原点を除くビット表現の上限
    bmub = 2^(n+m) - 1
-- DP配列
    cost :: Array Int [(Int, Double)]
    cost = listArray (0,bmub) $ map costf [0..bmub]
    costf 0  = [(goal, 0.0)]
    costf bm = [(k, costfk bm k) | k <- [0 .. n+m-1], testBit bm k]
    costfk bm k = minimum [c + dist j k * fact | (j, c) <- cost ! bm1]
      where
        bm1 = clearBit bm k
        fact = boost bm1
-- 最終結果
    goal = n + m
    ans = minimum
      [ costfk bm goal
      | mbm <- [0..2^m-1]
      , let bm = bit goal .|. (2^n-1) .|. shiftL mbm n
      ]

Arrayの内容を対応付けリスト [(Int, Double)] としたが、これを IntMap Double で実装した版の方が少し速くまたメモリも使わなかったのは、遅延評価によるスペースリークが原因だろうか？

F - Fishing

(2022-10-24) 解けたので追記。

問題 ABC274F

シグネチャを決める。第3引数は3タプルにするべきだが手抜きする。

abc274f :: Int      -- N
        -> Int      -- A
        -> [[Int]]  -- Wi,Xi,Vi
        -> Int      -- 答え
abc274f n a wxvs = ...

ある時刻 $t$ における最大値ならば、全ての魚の位置を求めて、網の位置をそこから $A$ だけ手前までのどこかの範囲にすることでその魚を捕まえることができるので、位置に対して捕まえられる魚の重さの和の関数の微分を作って累積和で積分して最大値をとる、というパターンが使える。これは時刻を固定して座標軸方向に関数を考えている。
しかし時刻も自分で決定する必要がある。時刻は離散的な整数でなく連続的な実数で要求されているので、魚の位置も実数で考える必要がある。

どれかの魚 $i$ が正解の網区間の最も左に居るとすると、これが $x$ の位置ぴったりにいると考えてよい。
この魚の位置と速度に対して他の魚 $j$ のそれを相対化して、$j$ が座標 $0$ から $A$ の範囲に滞在する時間範囲でスコアが $W_j$ だけ増える、という、時間軸方向に関数を考えて累積和から最大値を見つけることかできる。
「他の魚」の数が $N-1$、立ち上がりと立下りのエッジが $2(N-1)$ 箇所、これを Data.Map.fromListWith (+) で時刻に対してスコアの増減を記録した表を作るのに $O(N \log N)$ かかる。
魚 $i$ の選択を総当たりするため、結局 $O(N^2 \log N)$ かかる。$N \leq 2000$ と小さいので十分間に合いそう。

注意する必要があるのが、網の範囲が $x$ 以上 $x+A$ 以下 という点。普通に累積和を作るとき、座標が連続的かつ範囲の上限が「未満」ならば、微分関数は
$$
d(y) = \left \{\begin{array}{rl} W_j & y = xのとき \\ -W_j & y = x+Aのとき \\ 0 & それ以外のとき \end{array} \right .
$$
のようなスパイク関数で、$x+A$のときに負のスパイクを設定すればよい。
座標が離散的かつ範囲の上限が「以下」ならば、微分関数は
$$
d(y) = \left \{\begin{array}{rl} W_j & y = xのとき \\ -W_j & y = x+A+1のとき \\ 0 & それ以外のとき \end{array} \right .
$$
と負のスパイクをひとつ右にずらして対応できる。
しかし今回、座標が連続的かつ範囲の上限が「以下」で境界を含むので、前者で設定するとある魚が網を脱出する瞬間に他の魚が網に入るような場合に数え損なってしまう。後者でずらす量は無限小になって扱えない。

そこで、座標を (Ratio Int, Bool) とタプルで表現し、増やすスパイクは False に、減らすスパイクは True にすることで、無限小の隙間を作り出すことにする。

実装

魚のいずれか一つ $i$ を代表に選ぶ。これは群から抜き取らず、含めたまま計算すればよい。

import qualified Data.Map as M
import Data.Ratio

abc274f :: Int -> Int -> [[Int]] -> Int
abc274f n a wxvs = maximum $ map (score a wxvs) wxvs

score a wxvs (wi:xi:vi:_) = ...

ひとつの計算の中では

• 網の先頭の魚に対して他の魚を相対化し
• この相対化された魚が $[X_i,X_i+A]$ の範囲に入る瞬間の時刻、抜ける瞬間の時刻を求め、
• スパイク関数を足し合わせ
• 累積して最大値を求める

という手順でスコアを求める。

相対位置 $Y_j = X_j - X_i$, 相対速度 $U_j = V_j - V_i$ とおく。

score a wxvs (wi:xi:vi:) = ...
  where
    ...
      [ 
      | (wj:xj:vj:_) <- wxvs
      , let yj = xj - xi, let uj = vj - vi  -- 相対位置 Yj, 相対速度 Uj
      ...

網に入っている期間は $0 \leq Y_j + U_j \cdot t\leq A$ これを変形して
$U_j > 0$ のとき $(0 - Y_j)/U_j \leq t \leq (A - Y_j)/U_j$
$U_j < 0$ のとき $(A - Y_j) / U_j \leq t \leq (0 - Yj) / U_j$
$U_j = 0$ のとき、$0 \leq Y_j \leq A$ ならば常に範囲内、さもなくば常に範囲外

この時刻に重さ$W_j$ のスパイクを作る関数を立てる。
時刻は0以上のみ有意であり、負の時刻については全て集めてしまってよいので、これらの時刻が負であるときは $-1$ に集めることにする。
「常に」も $-\infty$ でなく $-1$ または $0$ に置けばよい。これは魚 $i$ 自身を含む。

spike yj uj wj =
  case compare uj 0 of
    GT -> [((t0, False), wj), ((ta, True), -wj)]
    EQ -> [((0 % 1, False), wj) | 0 <= yj, yj <= a]
    LT -> [((ta, False), wj), ((t0, True), -wj)]
  where
    m1 = -1 % 1
    t0 = let t = (0 - yj) % uj in if t < 0 then m1 else t
    ta = let t = (a - yj) % uj in if t < 0 then m1 else t

これを足し合わせて、時刻順に累積し、時刻が負の区間は捨て、0以上の区間の値の最大値をとる。

compute (wi,xi,vi) =
    maximum $ map snd $                   -- 値の最大値
    dropWhile ((0 >) . fst . fst) $       -- 時刻0以降のみを見て
    scanl1 add $                          -- 累積した結果に時刻を添えて
    M.assocs $                            -- 時刻の順に
    M.fromListWith (+) $                  -- 時刻ごとに集計
    [ tw
    | (wj,xj,vj) <- wxvs
    , tw <- spike (xj - xi) (vj - vi) wj  -- スパイクを生成 (Yj, Ujは名なしに)
    ]

add (_,acc) (t, w) = (t, acc + w)

常に含まれる値を時刻0に投入することで、他のスパイクが全て負の時刻に発生していたとしても、maximum の対象が空リストにならずに済むようにしている。

走らせてみると、結果は正しいが、時間が足りていない。えっ？

時刻のDouble化、無限小の時間差の表現の変更

マップのキーが複雑な構造をしていると、比較に時間がかかって不利なので、

• Ratio Int の精度をあきらめて Double で表す
• 無限小の隙間を時刻側で表現する( Map (Double, Bool) Int ）代わりに、時刻に対して足す値の累積と引く値の累積の二つを対応付ける（Map Double (Int,Int)）
  という改変を加える。
  scanl1 (+) で累積できなくなったので、その他の動作も含めて loop で全て処理する形に変更した。

abc274f n a wxvs = maximum $ map (score a wxvs) wxvs

score a wxvs (wi:xi:vi:_) =
    maximum $ loop 0 $ M.assocs $ M.fromListWith addp $
    [ tw
    | (wj:xj:vj:_) <- wxvs
    , tw <- spike (xj - xi) (vj - vi) wj
    ]
  where
    spike yj uj wj =
      case compare uj 0 of
        GT -> [(t0, (wj, 0)), (ta, (0,wj))]
        EQ -> [( 0, (wj, 0)) | 0 <= yj, yj <= a]
        LT -> [(ta, (wj, 0)), (t0, (0,wj))]
      where
        m1 = -1
        t0 = let t = (0 - yj) % uj in if t < 0 then m1 else t
        ta = let t = (a - yj) % uj in if t < 0 then m1 else t

addp (a,b) (c,d) = (a+c, b+d)

(%) :: Int -> Int -> Double
m % n = fromIntegral m / fromIntegral n

-- 増分wpを足した値だけ、時刻が負でないときだけ出力し、次の時刻に累積する
loop _ [] = []
loop acc ((t, (wp,wm)):twms)
  | t < 0     =        loop acc2 twms
  | otherwise = acc1 : loop acc2 twms
  where
    acc1 = acc  + wp
    acc2 = acc1 - wm

これでACできた。

時刻の Int 化

別のアプローチとして、パラメータの条件から、相対位置や相対速度は $[-10^4, 2 \times 10^4]$ の範囲にあり、時刻の除算を整数で充分な分解能と64ビット整数で桁あふれしない範囲で表現することが、分子を $10^9$ 倍することでできると、公式にリンクされたtoam氏による解説『Pythonで通す方法』に説明されていた。
このとき、Map より高速な IntMap が利用できる。また、無限小の時間差は $+1$ で表現することにする。

import qualified Data.IntMap as M

abc274f n a wxvs = maximum $ map (score a wxvs) wxvs

score a wxvs (wi:xi:vi:_) =
    maximum $ map snd $ dropWhile ((0 >) . fst) $
    scanl1 add $ M.assocs $ M.fromListWith (+) $
    [ tw
    | (wj:xj:vj:_) <- wxvs
    , tw <- spike (xj - xi) (vj - vi) wj
    ]
  where
    spike yj uj wj =
      case compare uj 0 of
        GT -> [(t0, wj), (succ ta, - wj)]
        EQ -> [( 0, wj) | 0 <= yj, yj <= a]
        LT -> [(ta, wj), (succ t0, - wj)]
      where
        m1 = -1
        t0 = let t = (0 - yj) % uj in if t < 0 then m1 else t
        ta = let t = (a - yj) % uj in if t < 0 then m1 else t

(%) :: Int -> Int -> Int
m % n = div (10^9 * m) n

add (_,acc) (t, w) = (t, acc + w)

より高速にACできた。
が、整数に写すところが危なっかしい（個人の感想、条件に依存して一般性に欠ける）ので好みではない。

感想

$O(N^2 \log N) \;; N \leq 2000$ がこんなに重いとは思わなかった。