ABC280 A～E をHaskellで

これまではおっとりと、コーディングと実行テストをAtCoder側の「コードテスト」でやっていたのを、手元のVSCodeでするように変えただけで、時間が少し稼げたような気がします。

A - Pawn on a Grid

問題 ABC280A

要は、$S_1$ ～ $S_H$ に出現する # の個数を数えろと言っている。

結果

main = getLine >> getContents >>= print . length . filter ('#' ==)

B - Inverse Prefix Sum

問題 ABC280B

シグネチャを決める。

abc280b :: Int    -- N
        -> [Int]  -- Si
        -> [Int]  -- 答え Ai

高校数学を思い出して
$$
\begin{array}{llrrl}
& A_1 + \dots + A_{k-1} & + \; A_k & = & S_k \\
-) & A_1 + \dots + A_{k-1} & & = & S_{k-1} \\
\hline
& & A_k & = & S_k - S_{k-1}
\end{array}
$$
ただし $S_0 = 0$ とする。

結果

abc280b n ss = zipWith (-) ss (0:ss)

C - Extra Character

問題 ABC280C

シグネチャを決める。

abc280c :: String  -- S
        -> String  -- T
        -> Int     -- 答え

前方から比較して、異なる文字が最初に現れた位置が答え（のひとつ）である。

結果

abc280e s t = succ $ length $ takeWhile id $ zipWith (==) s t

高速化

あまり考えずに上のコードで通してしまったが、$|S| \leq 5 \times 10^5$ ということで割と文字数があるので String でするのは無駄が多い。
ByteString でするものも作っておこう、と身構えたが、Data.ByteString.zipWith :: (Char -> Char -> a) -> ByteString -> ByteString -> [a] があるので何も変わらないようだ。

import qualified Data.ByteString.Char8 as BS

abc280e :: BS.ByteString -> BS.ByteString -> Int
abc280e s t = succ $ length $ takeWhile id $ BS.zipWith (==) s t

String版と比べて、時間 123ms→35ms、空間 48012KB→6600KB とかなり改善された。

D - Factorial and Multiple

問題 ABC280D

なかなか冷静になれなくて、3回もWAを提出してしまった。

シグネチャを決める。

abc280d :: Int  -- K
        -> Int  -- 答え

$K$ を素因数分解して $K = {p_1}^{n_1} \times {p_2}^{n_2} \times \dots \times {p_m}^{n_m}$ となったとする。
答え $N$ について、$N!$ はこれらを全て含む、つまり、1から$N$の素因数分解に $p_i$ は $n_i$ 個以上現れる。
階乗の性質から、$(A+1)!$ は $A!$ の倍数である。$A!$ の素因数は全て $(A+B)!$ に含まれる。$(B>0)$

$N$ の候補を $n$ とする。初期値は1。
$n$ を超える $p_i$ の倍数のうち最小のものを、次の候補 $n_1 = p_j \times a$ とする。
$n_1$ を素因数分解し、その中に含まれる $p_i$ の個数を $n_i$ からそれぞれ引く。0以下になった素因数は忘れる。
こうして、全ての素因数が出尽くしたときの $n$ の値が答えとなる。

結果

import qualified Data.IntMap as IM
import Data.List

abc280d :: Int -> Int
abc280d k = fst $ until (IM.null . snd) step (1, pfmap k)

step :: (Int, IM.IntMap Int) -> (Int, IM.IntMap Int)
step (c, m) = (c1, IM.differenceWith sub m cm)
  where
    c1 = minimum [p * divrup (succ c) p | p <- IM.keys m]
    cm = pfmap c1

pfmap :: Int -> IM.IntMap Int
pfmap = IM.fromAscList . map (\xs -> (head xs, length xs)) . group . primeFactors

sub a b
  | a > b = Just (a - b)
  | otherwise = Nothing

-- 切り上げ除算
-- @gotoki_no_joe
divrup x y = negate $ div (negate x) y

-- 素因数分解
-- @gotoki_no_joe
primeFactors :: Int -> [Int]
primeFactors n = loop n primes
  where
    primes = 2 : 3 : [y | x <- [5,11..], y <- [x, x + 2]]
    loop n pps@(p:ps)
      | n == 1    = []
      | n < p * p = [n]
      | r == 0    = p : loop q pps
      | otherwise = loop n ps
      where
        (q,r) = divMod n p

E - Critical Hit

問題 ABC280E

シグネチャを決める。

abc280e :: Int  -- N
        -> Int  -- P
        -> Int  -- 答え

スゴロク的なものの期待値は、ABC263Eに類題がある。263Eは進まずにループする場合もあって計算が面倒で、こちらの方が簡単。考え方は同じなのでなぞる。

攻撃で減らせる体力は1または2なので、残り1のときに2減らした-1が体力の下限。
$ex[-1 \leq i \leq N]$ を体力 $i$ から体力0以下にするための攻撃回数の期待値とする。
自明に、$ex[-1] = 0, ex[0] = 0$ はわかる。
体力 $i$ から、確率 $\frac{P}{100}$ で体力は $i-2$ になり、確率 $\frac{100-P}{100}$ で $i-1$ になる。よって、
$ex[i] = (ex[i-2] + 1) \times \frac{P}{100} + (ex[i-1] + 1) \times \frac{100-P}{100}$
$= \big \{ (ex[i-2] + 1) \times P + (ex[i-1] + 1) \times (100-P) \big \} / 100$
答えは $ex[N]$ である。
加減乗除を全て 998244353 を法とするモジュロ演算で行えばよい。

集めるDPなので、遅延配列で暗黙に計算できる。
直前とその前の項から次の項が得られるフィボナッチ数列のような漸化式として、明示的に前から確定させてもよい。

結果

遅延配列による暗黙の集めるDP実装

import Data.Array
import Data.List

abc280e :: Int -> Int -> Int
abc280e n p = ex ! n
  where
    pc = reg (100 - p)
    re100 = modRecip modBase 100
    ex = listArray (-1,n) $ map exf [-1..n]
    exf (-1) = 0
    exf 0 = 0
    exf i = mul re100 $ add 
            (mul p  $ add 1 $ ex ! (i-2))
            (mul pc $ add 1 $ ex ! pred i)

modBase = 998244353
reg x = mod x modBase
mul x y = reg (x * y)
add x y = reg (x + y)

-- @gotoki_no_joe
modRecip modBase a = powerish mul 1 a (modBase - 2)

-- @gotoki_no_joe
powerish mul i a b = foldl' {-'-} mul i [p | (True, p) <- zip bs ps]
  where
    bs = map odd $ takeWhile (0 <) $ iterate (flip div 2) b
    ps = iterate (\x -> mul x x) a

漸化式による方法

abc280e :: Int -> Int -> Int
abc280e n p = forthNth (succ n) exs
  where
    exs = 0 : 0 : zipWith step exs (tail exs)
    pc = reg (100 - p)
    re100 = modRecip modBase 100
    step ex2 ex1 = mul re100 $ add 
                   (mul p  $ add 1 $ ex2)
                   (mul pc $ add 1 $ ex1)

forthNth 0 (x:_ ) = x
forthNth n (x:xs) = x `seq` forthNth (pred n) xs

-- モジュロ演算は上と同じなので省略

前者が 73ms, 29276KB に対して、後者は 28ms, 4972KB であった。
計算の順序がはっきりしている場合には、その順に正格で評価を進める方が、特に空間効率を向上させる。

滑り込み

Dを2回WAして、後回しにしてEを解いてからDに戻って、さらにWAして絶望しかけ、思いついて最後の希望を提出したのが終了13秒前の 2022-12-03 22:39:47、判定中にタイムアップして、おそるおそる開いたら AC していた。

Eまでは実行時間が2桁ばかりというのも珍しいのでは。