ABC275 A～E+F をHaskellで

mod 998244353な回でした。
(2022/10/30 08:10) 問題Dのメモ化について追記。
(2022/10/30 20:00) 問題Fが解けたので追記。
(2022/10/31 05:00) 問題Fの説明を改善。

A - Find Takahashi

問題 ABC275A

シグネチャを決める。

abc275a :: Int    -- N
        -> [Int]  -- Hi
        -> Int    -- 答え

最大値を maximum で探す前に、その値が何番目の値かをタプルで追加しておけばよい。

結果

abc275a n hs = snd $ maximum $ zip hs [1..]

B - ABC-DEF

問題 ABC275B

シグネチャを決める。手抜きする。

abc275b :: [Int]  -- A～F
        -> Int    -- 答え

998244353は9桁の数、A～Fの上限 $10^{18}$ は18桁の数、64ビットな Int の maxBound = 9223372036854775807 は19桁の数で、最初の掛け算をする前に mod をとらないと桁あふれする恐れがある。

結果

abc275b args = (a `mul` b `mul` c) `sub` (d `mul` e `mul` f)
  where
    [a,b,c,d,e,f] = map reg args

modBase = 998244353
reg x = mod x modBase
mul x y = reg (x * y)
sub x y = reg (x - y)

Integer で計算するなら何も考えなくてもできてしまう。

abc275b :: [Integer] -> Integer
abc275b [a,b,c,d,e,f] = mod (a * b * c - d * e * f) 998244353

C - Counting Squares

問題 ABC275C

シグネチャを決める。

abc275c :: [String]  -- Si
        -> Int       -- 答え

例を見ればわかるが、正方形はまっすぐに置いてあるとは限らない。
点 $A,B,C,D$ が正方形をなすとする。$A$ から見た $B$ の相対位置が「右に $dx$ 下に $dy$」であるとき、$A$ から見た $D$ の相対位置は「下に $dx$ 左に $dy$」、$D$ から見た $C$ の相対位置は「右に $dx$ 下に $dy$」となる。
よって、全ての # の位置から任意に2つ選んだ組み合わせを $A,B$ としたとき、$C,D$ に相当する位置に # があるかを調べ、そのような選択の個数を数える。
$dx > 0, dy \geq 0$ の場合に限定することで、重複を避けることができる。
あとそもそも、A,B と割り当てたら B,A と割り当てて数えなおす必要はない。

結果

画面を配列で管理すると、添え字の外に出た場合を気にするのが面倒なので、座標の集合で扱う。

import qualified Data.Set as S

abc275c :: [String] -> Int
abc275c css = cnt
  where
    ijs = [(i,j) | (cs, i) <- zip css [0..], ('#', j) <- zip cs [0..]]
    ijS = S.fromAscList ijs
    cnt = length
      [ ()
      | (a,b):cds <- tails ijs, (c,d) <- cds
      , let dx = c - a, let dy = d - b
      , dx > 0, dy >= 0
      , S.member (a - dy, b + dx) ijS
      , S.member (a - dy + dx, b + dx + dy) ijS
      ]

D - Yet Another Recursive Function

問題 ABC275D

シグネチャを決める。

abc275d :: Int  -- N
        -> Int  -- 答え

安直にメモ化でやろうとすると、$N \leq 10^{18}$ は大きすぎてメモ帳が用意できない。
フィボナッチ数列ならすぐ手前の数だけから計算できるため容易にループで計算できるが、この関数はかなり離れた位置の値を必要とするのでそれもできない。

最終的に到達する基底部が $f(0) = 1$ であることに注目する。$N$ から始めて再帰的に $f(k) = f(\lfloor k/2 \rfloor) + f(\lfloor k / 3 \rfloor)$ で降りていくとき、「今の引数の結果が最終結果に何個足しこまれるか」という個数を引数それぞれに対して考える。例えば $f(12)$ を計算するとき

12が1回　12を展開して
12/2=6が1回、12/3=4が1回　6を展開して
6/2=3が1回、6/3=2が1回、4が1回　4を展開して
3が1回、2が1回、4/2=2が1回、4/3=1が1回 ＝ 3が1回、2が2回、1が1回　3を展開して
3/2=1が1回、3/3=1が1回、2が2回、1が1回 ＝ 2が2回、1が3回　2を展開して
2/2=1が2回、2/3=0が2回、1が3回 ＝ 1が5回、0が2回　1を展開して
1/2=0が5回、1/3=0が5回、0が2回 ＝ 0が12回 ＝ 12f(0) = 12

となる。「0がx回」のみになったときのxが答えである。

結果

import qualified Data.IntMap as IM

abc275d :: Int -> Int
abc275d n = im IM.! 0
  where
    im = until done step $ IM.singleton n 1

done im = IM.size im == 1 && IM.member 0 im

step im = im3
  where
    ((k,v), im1) = IM.deleteFindMax im
    im2 = IM.insertWith (+) (div k 2) v im1
    im3 = IM.insertWith (+) (div k 3) v im2

追記：メモ化

朝になったらTwitterのトレンドに「メモ化再帰」(とPython)が上がってて、ABC275Dが、実際に呼ばれたところのメモをとるだけで解ける、Pythonにはそういうライブラリがある、ということらしい。
Haskellのメモ化は「可能性のある全ての引数に対する結果を遅延評価で書き込んだ配列を作る」スタイルは簡単なのだけど、「実際に呼ばれたところ」の結果をピンポイントに記録するには mutable map をメモ帳として使う必要があって、そのメモ帳を次の呼び出しに渡すようにするのが面倒くさい。が、わかりやすい教科書再帰ドリル(7)：メモ化 - Haskellでのメモ化があるので、これに従ってやってみよう。

import qualified Data.IntMap as IM

main = readLn >>= print . abc275d

abc275d :: Int -> Int
abc275d = fst . f IM.empty

f :: IM.IntMap Int -> Int -> (Int, IM.IntMap Int)
f im 0 = (1, im)
f im n =
  case IM.lookup n im of
    Just r  -> (r, im)
    Nothing -> (r, im3)
  where
    (a, im1) = f im  (div n 2)
    (b, im2) = f im1 (div n 3)
    r = a + b
    im3 = IM.insert n r im2

通った。

教科書の続き再帰ドリル(7)：メモ化 - memoize パッケージのmemoizeに、これを一般化したものがある、とあるのでそれもやってみよう。

import Data.Function.Memoize

main = readLn >>= print . abc275d

abc275d :: Int -> Int
abc275d = f_memo

f_memo = memoize f

f :: Int -> Int
f 0 = 1
f n = f_memo (div n 2) + f_memo (div n 3)

めっちゃコードがシンプルになったが、このモジュールはAtCoderには入っていなかった。

Could not find module ‘Data.Function.Memoize’

これの中身はTemplate Haskellを使ってごりごりしているっぽいので、核を抜き出してくるのも大変そう。
中の人の解説 Haskell memoization を真似してみよう。目標の関数をモナド化する必要は残るが、メモ動作そのものは隠蔽される。
またこれは、複数回の「本命の」呼び出しの間でメモが共有されないので、この問題については支障がないが、もったいない。

import qualified Data.IntMap as IM
import Control.Monad.State
import Data.Maybe

main = readLn >>= print . abc275d

abc275d :: Int -> Int
abc275d = memoizeM f

-- 自分を最初の引数にとり、モナド化した目標関数
f _  0 = return 1
f mf n = do
  a <- mf (div n 2)
  b <- mf (div n 3)
  return $ a + b

-- https://gist.github.com/Janiczek/3718805

type StateMap = State (IM.IntMap Int) Int

memoizeM :: ((Int -> StateMap) -> (Int -> StateMap)) -> (Int -> Int)
memoizeM t = \x -> evalState (f x) IM.empty
  where
    f x = do
      im <- get
      case IM.lookup x im of
-- メモに答えがあればそれを返す
        Just r -> return r
-- なければ計算する
        Nothing -> do
-- 目標関数 t に、「それ自身」としてメモ機能を追加した f を渡して実行
          y <- t f x
-- 結果をメモに追記
          im <- get              -- (※)
          put (IM.insert x y im)
          return y

ACした。

(※)の get をサボって im を流用すると、再帰呼び出しで追記した部分をなくしてしまうのでまちがい。

E - Sugoroku 4

シグネチャを決める。

abc275e :: Int  -- N
        -> Int  -- M
        -> Int  -- K
        -> Int  -- 答え

第1ターンから第 $K$ ターンまで、マス $0$ からマス $N$ までの位置にコマが存在する場合の数をまっすぐ生成して数える。
初期状態（第0ターン）は $(1,0,\dots,0)$ である。
前のターンの結果が $(C_0,\dots,C_{N-1},0)$ のとき、添え字を $1$～$M$ だけ足した位置にそれらを足し込むことで次のターンの状態を作れる。ただし$N$を超えた分は折り返す。
それぞれのターンにおけるマス $N$ の場合の数を取り出す。

第 $p$ ターンにゴールする場合の数が $q$ であるとき、それはゴールできる確率を $\frac{q}
{M^p}$ 増やす。

ステップ1回の計算量が $O(NM)$、これを $K$ 回繰り返すので $O(NMK)$ で、$N \leq 1000, M \leq 10, K \leq 1000$ なので $10^7$ ぐらいなので普通に間に合う。

結果

import Data.Array.Unboxed
import Data.List

abc275e :: Int -> Int -> Int -> Int
abc275e n m k = ans
  where
    initial = accumArray (+) 0 (0,n) [(0,1)] :: UArray Int Int
    step arr = accumArray add 0 (0,n)
      [ (if j <= n then j else (n + n - j), a)
-- 0からN-1までを足しこめばいい マスNからは進まない
      | (i,a) <- zip [0..pred n] (elems arr)
      , j <- [i+1..i+m]]
    cnts = map (! n) $ take (succ k) $ iterate step initial
    recipM = modRecip modBase m
    ans = foldl1 add $ zipWith mul cnts $ iterate (mul recipM) 1

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)
  where
    mul x y = mod (x * y) modBase

-- @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

F - Erase Subarrays

(2022/10/30 20:00 解けたので追記)

問題 ABC275F

シグネチャを決める。

abc275f :: Int    -- N
        -> Int    -- M
        -> [Int]  -- Ai
        -> [Int]  -- 答え

初めのアプローチ

数列の部分区間に関して、1～Mの数それぞれを作るのに必要な操作回数が得られているとする。
隣接する二つの区間を繋ぎ合わせた区間について、それぞれの数を作るのに必要な操作回数をこれから導けるだろうか。
区間$S_1$で数$a$を作るのに必要な操作回数を$b$、区間$S_2$で数$c$を作るのに必要な操作回数を$d$とするとき、区間$S_1S_2$で数$a+b$は操作回数$b+d$で作れると言えるだろうか。その限りではない。$S_1$で$a$を作るやり方か$S_1$の末尾の項を削除しており、$S_2$で$c$を作るやり方が先頭の項を削除しているとき、両者にまたがる区間で操作できて一回省略できる。
よって、区間に対して、1～Mの数をそれぞれ作るのに必要な操作回数、ただし、その操作が両端の項を削除する／しないの4とおりそれぞれについて別に数える、という情報を持っていれば、それらをつなぎ合わせることができる。
初期値はそれぞれ長さ1の部分列$[A_i]$について、両端とも残して$A_i$を操作0回、両端とも削除して0を操作1回、それ以外は不可能、とする。
ボトムアップのマージソートのように、隣接する区間どうしを繋ぎ合わせることを繰り返し、数列全体について求められたら、両端の残し方を無視してそれぞれの値について最小の回数を決定する。

しかしこの方法は間に合わなかった。あとREが1回出る。

つなぎ方を変える

長さ1の区間N個を、1,3,5...番目をそれぞれ2,4,6...番目とつなぎ合わせる、と完全二分木の形でやる代わりに、単純に前から1つずつ繋いでいってもよいはず。それは単なる foldl になる。

すると、前側の列の左端が使用済みかどうかは使わない情報になる。末尾を使ったか使わなかったかだけでよい。後ろ側の列は長さ1なので、両端というかそれを使うか使わないかだけとなる。結局、両端に関してそれぞれ4通りの場合分けをする代わりに、前回 $A_{i-1}$ を使ったか、今回 $A_i$ を使うかのそれぞれ2択だけでよくなる。

前から $i$ 項の部分列について、第 $i$ 項は使って、和 $v$ を作るために必要な操作の回数を記録する。これを $C_i[0 \leq v \leq M]$ とする。
前から $i$ 項の部分列について、第 $i$ 項は使わずに、和 $v$ を作る操作の回数を記録する。これを $D_i[0 \leq v \leq M]$ とする。
（一般的には、三次元配列 $dp[0 \leq i \leq N][0 \leq v \leq M][Bool]$ で、先頭から現在の項までで、1からMの値を作るための最小操作回数、最後の項を使ったか否かがBool値の添え字、とするところ。）

この二つの配列のペア $(C_i, D_i)$ と $A_{i+1}$ から次の配列のペア $(C_{i+1},D_{i+1})$ を作る漸化式（つまり foldl のステップ関数）を考える。

$A_{i+1}$ を使わないで値 $v$ を作る方法の最小回数は

• $A_i$は使って作る回数 $C_i[v]$ から、使わない回数を+1
• $A_i$も使わないで作る回数 $D_i[v]$

の小さい方なので、$D_{i+1}[v] = \min(C_i[v] + 1,D_i[v])$ となる。

$A_{i+1}$ を使うとき、この数が足しこまれるので、値 $v$ を作る方法の最小回数から、値 $v+A_i$ を作る方法の最小回数が導かれて、

• $A_i$は使って作る回数 $C_i[v]$
• $A_i$は使わないで作る回数 $D_i[v]$

の小さい方なので、$C_{i+1}[v+A_i] = \min(C_i[v],D_i[v])$ となる。

初期値は $C_0[0] = 0, C_0[v>0] = \infty, D_0[v] = \infty$ でよい。
最終結果は、それぞれ $C$ と $D$ の小さい方を選ぶ。（どちらも $\infty$ なら $-1$）

漸化式を見ると、添え字について使用が前後していないので、配列を使わずに $C,D$ はリストで充分扱える。すなわち、$C,D$ を cee, dee とすると次のようにできる。$C$ の添え字をずらすための埋め草は、次の周回で succ される可能性があるため maxBound では大きすぎる。

tooBig = div maxBound 2

step (cee,dee) ai = (cee1, dee1)
  where
    dee1 = zipWith min dee $ map succ cee
    cee1 = take (succ m) (replicate ai tooBig ++ zipWith min cee dee)

ただし本当にこのままだと、リストの内容の計算が遅延評価で無駄に後回しにされて TLE するので、評価を強制する必要がある。

結果

import Data.List
import Control.DeepSeq

abc275f :: Int -> Int -> [Int] -> [Int]
abc275f n m as =
    [ if r >= tooBig then -1 else r
    | r <- tail $ zipWith min ceeN deeN]
  where
    m1 = succ m
    initial = (0 : replicate m tooBig, replicate m1 tooBig)  -- cee0, dee0
    (ceeN, deeN) = foldl' step initial as
    step (cee, dee) ai = force (cee1, dee1)
      where
        cee1 = take m1 $ replicate ai tooBig ++ zipWith min cee dee
        dee1 = zipWith min dee $ map succ cee

tooBig :: Int
tooBig = div maxBound 2

force を忘れるとTLEするのが、Haskellの本当にもったいないところ。