paiza「島探し」と「十億連勝」(paizaランク S 相当)」もHaskellで

残り二つも解けたので書いておく。

島探し

問題

シグネチャを決める。

solve :: Int -- M
      -> Int -- N
      -> [[Int]] -- 0または1
      -> Int -- 答え

自分の解答

• Union-Findを用いる
• 黒いマスに順に番号を振り、$(M+1)\times(N+1)$の表にその番号を書いておく。（白いマスは0とする）
• それぞれの黒いマスについて、上隣または左隣のマスも黒いとき、その番号とUnionする
• 黒いマスの個数ーUnionした回数、が答え

import Data.Array
import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad.ST

solve :: Int -> Int -> [[Int]] -> Int
solve m n ass = runST $
  do
    uf <- newUF (succ k)
    foldM (\cnt (a,b) -> do
      r <- uniteUF uf a b
      return $ if r then pred cnt else cnt
      ) k neighbors
  where
    k = sum $ concat ass -- 1の総数
    arr = accumArray (+) 0 ((0,0),(n, m)) $
          flip zip [1..]
          [(i,j) | (i,as) <- zip [1..] ass, (j,1) <- zip [1..] as]
    neighbors = concat
      [ [(a,b) | let b = arr ! (pred i,j), b > 0] ++
        [(a,c) | let c = arr ! (i,pred j), c > 0]
      | i <- [1..n], j <- [1..m], let a = arr ! (i,j), a > 0 ]

実際にはこれに Union-Findのコードも加わる。

公式解説のやり方

PAINT方式でコンパクトなコードになっていた。写経する。

import Data.Array.Unboxed

solve :: Int -> Int -> [[Int]] -> Int
solve m n ass = ans
  where
    arr0 :: UArray (Int,Int) Bool
    arr0 = accumArray (flip const) False ((0,0),(succ n,succ m))
           [((i,j),True) | (i,as) <- zip [1..] ass, (j,1) <- zip [1..] as]
    (ans, _) = foldl outer (0, arr0) [(i,j) | i <- [1..n], j <- [1..m]]
    outer st@(cnt, arr) ij
      | arr ! ij  = (succ cnt, inner arr [ij])
      | otherwise = st
    inner :: UArray (Int,Int) Bool -> [(Int,Int)] -> UArray (Int,Int) Bool -- paizaが要求
    inner arr [] = arr
    inner arr (ij@(i,j):ijs)
      | arr ! ij = inner (arr // [(ij,False)]) $ (pred i,j) : (succ i,j) : (i,pred j) : (i,succ j) : ijs
      | otherwise = inner arr ijs

• arr0 マス目を白黒に塗っておく。
• foldl outer 全てのマスについて、そこが黒いとき
• inner 隣接するマスも全て白くする。
• inner を起動した回数が答え

immutable arrayの更新で手抜きしたが、全然間に合ったのは Bool 配列だったからだろうか？
innerの型シグネチャは、paizaの処理系が要求したので入れている。

十億連勝

問題

シグネチャを決める。

solve :: Int -- N
      -> Int -- X
      -> [Int] -- Ai
      -> Int -- 答え

なんだかややこしいけれど、

• ステージはいくつかの試合からなる
• ステージが終了するのは、全勝するか、途中で最初に負けたそのとき
• ステージを跨ぐものも考慮して、最大の連勝数がXである場合の数を数える

ということ。

自分の解答

パラメータの規模感からして、何らかのDPで解けるのだろうけど、どんな状態を追跡したらよいのだろう。
「これまで最大の連勝数（現在進行形のものを除く）」と「現在の連勝数」の対を状態としたDPでできないだろうか。考えるべき状態数が多すぎて間に合わないような気もするけれど。

直前のステージ $i-1$ を「(最大 $p$ 連勝, 現在 $q$ 連勝中)」という状態で終えたとき、
ステージ $i$ は最大で $A_i$ 試合あり、

• $0 \leq a < A_i$ 勝して次の試合で負け、現在の連勝記録が停止したとき、状態 $(\max(p, q+a), 0)$ に遷移
• 全勝して連勝を継続させたとき、$(p, q + a)$ に遷移

となるので、それらの場合の数を追跡する。
初期状態は $(0,0)$ が1とおり、である。
$p, q \leq X$ な場合のみ数えればよい。
最終ステージが終了した後、$q$ の方も合わせて、$(X, *)$ と $(*, X)$ の場合の数の総和が答え。

import qualified Data.Map as M

solve :: Int -> Int -> [Int] -> Int
solve n x as = add 0 $ sum [r | ((p,q),r) <- M.assocs mN, max p q == x]
  where
    mN = foldl' step (M.fromList [((0,0),1)]) as
    step m ai = M.fromListWith add
      [ ((p1, q1), r)
      | ((p,q),r) <- M.assocs m
      , a <- [0 .. ai]
      , let p1 = max p $ q + a, p1 <= x
      , let q1 = if a == ai then q + a else 0, q1 <= x
      ]

add :: Int -> Int -> Int
add x y = mod (x + y) (10^9)

状態対が $X^2 = 10^{18}$ とおり、さらにステージでの勝利数が毎回 $(A_i = 10^9) + 1$ 通りを $N \leq 4000$ 回するのでは到底間に合いそうになかったが、間に合った。テストケース優しすぎか。

ともかくこれで公式解説が読める。

公式解説のやり方

上の「最大 $p$ 連勝」を $0 \leq p \leq X$ について分けて考える必要は全くなくて、「ちょうど $X$ 連勝したことがあるかどうか」の Bool を追いかけるだけでよい。
言われてみればそのとおり。
ただし、$q + A_i$ が $X$ を越えうるかどうかで、どの場合にいくつ足し込めばよいかという推移関係がややこしい。解説にある式が正しいことは追って理解できたが、いきなりこれを立式できる気がしない。

solve :: Int -> Int -> [Int] -> Int
solve n x as = summ [c | (wb,c) <- M.assocs mN, prop wb]
  where
    mN = foldl' step (M.fromList [((0,False),1)]) as
    step m ai = M.fromListWith add $ concat
      [ if w + ai > x then [((0,True), c), ((0, b), mul c $ x - w)] else [((w + ai, b),c),((0, b), mul c ai)]
      | ((w,b),c) <- M.assocs m
      ]
    prop (w,b) = b || w == x

reg :: Int -> Int
reg x = mod x (10^9)

add :: Int -> Int -> Int
add x y = reg $ x + y

mul :: Int -> Int -> Int
mul x y = reg $ x * y

summ :: [Int] -> Int
summ = foldl' add 0

公式解説に合わせたので、変数名と、対の左右が上と食い違っている。

終わりに

スキルチェックでHaskell使えるようにしてほすい…