ABC298 A～G をHaskellで

トヨタが冠なのはいつも難しいですね。全然解けませんでした。

A - Job Interview

問題 ABC298A

シグネチャを決める。

abc298a :: Int     -- N
        -> String  -- S
        -> Bool    -- 答え

二つの条件を翻訳する。

• 「良」と評価した面接官が少なくとも 1 人いる ＝ 'o' が $S$ に含まれる
• 「不可」と評価した面接官がいない ＝ `'x' が $S$ に含まれない

結果

abc298a n s = elem 'o' s && notElem 'x' s

試しにやってみたところ、BingAIさんも完璧な答えを返しました。

B - Coloring Matrix

問題 ABC298B

シグネチャを決める。

abc298b :: Int      -- N
        -> [[Int]]  -- Aij
        -> [[Int]]  -- Bij
        -> Bool     -- 答え

回転は、問題文にあるとおりにやらなくても、transpose と reverse でできる。
（回転の向きが逆かもしれないけど、全て試すので問題ない。）

結果

import Data.List

abc298b :: Int -> [[Int]] -> [[Int]] -> Bool
abc298b n ass bss = any (prop bss) $ take 4 $ iterate rot ass
 
prop bss ass = null [() | (bs,as) <- zip bss ass, (0,1) <- zip bs as]
 
rot = reverse . transpose

こちらもBingAIさんに聞いてみたら、「$A_{i,j} = 1$ ならば $B_{i,j} = 1$」を「$A = B$」と勘違いした（以外は完璧な）答えでした。
実はその間違いを自分も最初やらかしたので、非常に人間臭い間違いをする機械だなぁという。

C - Cards Query Problem

問題 ABC298C

シグネチャを決める。
クエリによってデータの更新が起きるので、mutable array を使って計算する。
出力も内部でやってしまい、IO モナドで全体を回す。

abc298c :: Int      -- N
        -> Int      -- Q
        -> [[Int]]  -- queryi
        -> IO ()

import Data.Array.IO

クエリ 2 i に対して、箱 i の内容を、数の重複も含めて答える必要があるので、箱ごとにその内容の多重集合が必要。数ごとの要素数を持つ IntMap で表現する。

import qualified Data.IntMap as IM

クエリ 3 k に対して、数 k が入っている箱の番号を全て答える必要があるので、数ごとにそれが入った箱の番号の集合が必要。重複は無視するので IntSet で表現する。

import qualified Data.IntSet as IS

つまり、クエリ 1 k i に対して、箱の内容配列と数を持つ箱配列の両方を更新する必要がある。

結果

abc298c n q qs = do
  box <- newArray (1,n) IM.empty       -- 箱の内容配列
  num <- newArray (1,200000) IS.empty  -- 数を持つ箱配列
  forM qs $ qi -> do
    case qi of
      [1, k, i] -> q1 box num k i
      [2, i]    -> readArray box i >>= q2
      [3, k]    -> readArray num k >>= q3

q1 :: IOArray Int (IM.IntMap Int)
   -> IOArray Int (IS.IntSet)
   -> Int -> Int -> IO ()
q1 box num k i = do
  writeArray box i . IM.insertWith (+) k 1 =<< readArray box i
  writeArray num k . IS.insert i           =<< readArray num k
 
q2 ms = putStrLn $ unwords [show i | (i,m) <- IM.assocs ms, _ <- [1..m]]
 
q3 = putStrLn . unwords . map show . IS.elems

D - Writing a Numeral

問題 ABC298D

シグネチャを決める。
またクエリ問題。
mutable array は必要なさそうなので、pureに計算する。

abc298d :: Int      -- Q
        -> [[Int]]  -- queryi
        -> [Int]    -- 答え

現在の $S$ を10進表記として解釈した値 $v$ を $\bmod 998244353$ で保持する。

• クエリ 1 x に対しては、$10v + x$ に更新する。
• クエリ 3 に対して、$v$ を出力する。

クエリ 2 に対して、最も古い数字を、最上位桁から取り除く必要がある。
このため、$v$ だけでなく、入力された数字列に FIFO でアクセスできる必要があるので、 Data.Sequence に保持する。
最も古い数字が $d$、現在の $S$ の長さが $k$ のとき、$d \times 10^{k-1}$ を $v$ から引くことで、最上位桁を取り除ける。これも $\bmod 998244353$ で計算すればよい。
$10^{k-1} \bmod 998244353$ は $0 \leq k \leq Q-1$ について持っておけばよい。

結果

import Data.Array
import qualified Data.Sequence as Q

abc298d :: Int -> [[Int]] -> [Int]
abc298d q qs = loop (Q.singleton 1) 1 qs
  where
    loop :: Q.Seq Int -- 現在の S
         -> Int       -- 現在の値 v
         -> [[Int]]   -- queryi
         -> [Int]     -- 答え
    loop s v         ([1,x]:qs) = loop (s Q.|> x) (reg $ 10 * v + x) qs
    loop (d Q.:<| s) v ([2]:qs) = loop s (reg $ v - d * base ! Q.length q) qs
    loop s v           ([3]:qs) = v : loop s v qs
    loop _ _ []                 = []

    base = listArray (0, pred q) $ take q $ iterate (reg . (10 *)) 1

reg x = mod x 998244353

E - Unfair Sugoroku

問題 ABC298E

これ以降は判らなかったので、解説を見てやっていきます。

シグネチャを決める。

abc298e :: Int -> Int -> Int -> Int -> Int -> Int
abc298e n a b p q = ...

「すごろくは後ろから」
各状況での高橋君の勝率をDPで求める。高橋君の位置を$i$、青木君の位置を$j$とする。現在、二人のどちらの手番かの変数を$t$とし、高橋君のとき0、青木君のとき1をとるとする。
これらを変数に持つ $prop[i, j, t]$ を考える。
要求されている答えは $prop[A,B,0]$
高橋君がゴール済みの状況での勝率は1なので $prop[N,j<N,t] = 1$
青木君がゴール済みの状況での勝率は0なので $prop[i<N,N,t] = 0$
高橋君の番 $prop[i<N,j<N,0]$ の勝率は、$1 \leq k \leq P$ に等確率で進んだ先の勝率を、進む確率を掛けて足し合わせたものになるので $prop[i,j,0] = \big (\sum_{k=1}^P prop[i+k, j, 1] \big ) / P$
青木君の番 $prop[i<N,j<N,1]$ の勝率は、$1 \leq k \leq Q$ に等確率で進んだ先の勝率を、進む確率を掛けて足し合わせたものになるので $prop[i,j,1] = \big (\sum_{k=1}^Q prop[i, j+k, 0] \big ) / Q$
ここで、ゴールを超えて進まないように、添え字が$N$を超えるないように $\min(N, i+k)$ のようにクランプしておく必要がある。

モジュロ演算での確率は、割り算がモジュロ逆数を掛けることで行われる以外、普通の計算と同じにして構わない。

結果

いつもの、遅延配列を用いた集めるDP。

import Data.Array
import Data.List

abc298e :: Int -> Int -> Int -> Int -> Int -> Int
abc298e n a b p q = prop ! (a,b,0)
  where
    recipP = modRecip p   -- モジュロ逆数 1/P
    recipQ = modRecip q   -- モジュロ逆数 1/Q
    bnds = ((a,b,0),(n,n,1))
    prop = listArray bnds $ map f $ range bnds
    f (i,j,_) | i == n = 1
              | j == n = 0
    f (i,j,0) = mul recipP $ summ [prop ! (nlim $ i+k, j, 1) | k <- [1..p]]
    f (i,j,1) = mul recipQ $ summ [prop ! (i, nlim $ j+k, 0) | k <- [1..q]]
    nlim = min n

-- モジュロ演算
modBase = 998244353
reg x   = mod x modBase
mul x y = reg $ x * y
summ    = foldl1' add
add x y = reg $ x + y

-- モジュロ逆数
modRecip a = powerish mul 1 a (modBase - 2)

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 - Rook Score

問題 ABC298F

これも解説を見て実装。途中で何となく既視感が。

シグネチャを決める。$r_i,c_i,x_i$は手抜きする。

abc298f :: Int      -- N
        -> [[Int]]  -- ri,ci,xi
        -> Int      -- 答え

数が一つは書かれた行や列だけを考えればよい。
行ごと、列ごとに、数の総和をとる。これらの値を「行和」「列和」と呼ぶことにする。
それぞれの行について、列との交点は、行和＋列和 になるので、列和が最大の列を選べばよい。

ところが、その交点に数が書いてあったら、その数を引く必要があるので、列和の大きい順で二番手もチェックして、大きい方を選ぶ必要がある。
ところが、その交点にも数が書いてあったら…と続く。
列和の大きい順に調べて、数の書いていない交点に遭遇したとき、それより列和の小さい列が最大値をもたらすことはないので、そこで打ち止めることができる。

という調査を全ての行に関して行い、それらから最大値を選べばよい。
このとき、列和の大きい順で、次点を探さないといけなくなる回数は、数を書いた回数の $N$ に等しいので、全体で最大 $N+1$ 箇所だけ調べることになる。

結果

import qualified Data.IntMap as IM
import Data.Tuple
import Data.Maybe

abc298f :: Int -> [[Int]] -> Int
abc298f n rcxs = maximum $ map getCands $ IM.assocs rows
  where
-- 行和を持つIntMap
    rows = IM.fromListWith (+) [(r,x) | r:_:x:_ <- rcxs]
-- 列和を大きい順に、(値, 列番号)
    cols = sortBy (flip compare) $ map swap $ IM.assocs $
           IM.fromListWith (+) [(c,x) | _:c:x:_ <- rcxs]
-- 行に対して、数を書いた列番号とその値の対リスト、のIntMap
    ps = IM.fromListWith (++) [(r,[(c,x)]) | r:c:x:_ <- rcxs]
-- 行r（その行和はrv）の交点の最大値を求める
    getCands (r, rv) = rv + maximum (take 1 (map negate bs) ++ as) -- ※4
      where
        psr = IM.fromList $ ps IM.! r
        (as,bs) = span (0 <=)                                   -- ※3
                  [ maybe (negate cv) (cv -) $ IM.lookup c psr  -- ※2
                  | (cv, c) <- cols]                            -- ※1

getCands の中では、
※1：列和の大きい順に調べ、
※2：その行その列に数が書いてあったら、列和からその値を引いた値、数が書かれていなければ列和のマイナス（打ち切りのマーカー）を吐き出し、
※3：0以上（数が書いてあったマスの値）をas、負数（数が書いてなかったマスの値）をbsに分けて
※4：as全体と、bsをプラスに戻したものの先頭（ただし空の可能性あり）から最大値を選び、行和と足す
とやっている。符号でなく Either を使う方が「正しい」という意味では手抜き。

G - Strawberry War

問題 ABC298G

これも解説を見て、もまだわからなかったのでアライさんのPython版をなぞっただけ。

考える

import Data.Array
import Data.List
import qualified Data.IntSet as IS

abc298g :: Int      -- H
        -> Int      -- W
        -> Int      -- T
        -> [[Int]]  -- sij
        -> Int      -- 答え
abc298g h w t sss = ...
  where

ケーキの区画で数える代わりに、カットする碁盤の目に1から番号を付けて呼ぶ。
線1番から線2番の間を取り出すと、つまり区画1だけを指している、という感じ。
（上開区間）

色々な区画の繋がりについてイチゴの数を高速に数えるために、二次元累積和を使う。
berries u d l r は、線uからd、lからrまでの区画のイチゴの数を返す。

    t1 = succ t
    h1 = succ h
    w1 = succ w
    accs = listArray ((1,1),(h1,w1)) $ concat $
           scanl (zipWith (+)) (replicate w1 0) $
           map (scanl (+) 0) sss
    berries u d l r = accs ! (d,r) + accs ! (u,l) - accs ! (d,l) - accs ! (u,r)

下限として可能性のある数は、あらゆるberriesの結果の値。

    lowers = IS.fromList
      [ berries u d l r
      | u <- [1..h], d <- [succ u..h1]
      , l <- [1..w], r <- [succ l..w1]
      ]

求める答えを与える下限 lower をこれで総当たりして試す。
solve lower 関数は、下限を lower としたときの、最大値の最小値を返す。
ただしそのような切り分け方ができない場合は -1 を返す。
という関数があったとすると、lower を下回らないような結果の最小値が答えになる。

abc298g h w t sss = minimum
    [ res - lower
    | lower <- IS.elems lowers
    , let res = solve lower
    , res >= lower
    ]

肝心の関数 solve は、
「線uからd、lからrまでをc個に分割したときの最大値、ただし下限lowerを下回る場合は-1」
が全て入った配列arr ! (u,d,l,r,c) があるとき、arr ! (1,h1,1,w1,t1) を返せばよい。
この配列は集めるDPで作れる。

    solve lower = arr ! (1,h1,1,w1,t1)
      where
        bnds = ((1,2,1,2,1),(h,h1,w,w1,t1))
        arr = listArray bnds $ map f $ range bnds
-- 配列の内容を計算する関数
        f (u,d,l,r,c)
          | c == 1    = if b < lower then -1 else b -- 分割1個ならもう分割しない
          | null recs = -1                          -- 分割失敗のとき -1
          | otherwise = minimum recs                -- 成功する分割の最小値
          where
            b = berries u d l r
            recs =
              [ max a1 a2
              | m <- [succ u..pred d]                  -- u-d方向に二分して総当たり
              , cc <- [1..pred c]                      -- 分割数も分けて全て試す
              , let a1 = arr ! (u, m, l, r, cc)
              , a1 >= lower                            -- 分割に失敗したら放棄
              , let a2 = arr ! (m, d, l, r, c - cc)
              , a2 >= lower
              ] ++
              [ max a1 a2
              | m <- [succ l..pred r]                  -- l-r方向に二分して総当たり
              , cc <- [1..pred c]                      -- 以下、上に同じ
              , let a1 = arr ! (u, d, l, m, cc)
              , a1 >= lower
              , let a2 = arr ! (u, d, m, r, c - cc)
              , a2 >= lower
              ]

分割に失敗した場合にリスト内包表記で打ち切りされることで、枝刈りは実現されていると思っている。

性能比較

解	時間(ms)	メモリ(MB)
これ	5700	25
手本にしたアライさんのPython	3800	90
なんかすごいHaskell解	25	7

時間も空間もなんか桁を間違えているような結果があって気が遠くなる。