ABC363 A～F をHaskellで

今回はなんかメロメロでした。

A - Piling Up

問題 ABC363A

シグネチャを決める。

abc363a :: Int -- R
        -> Int -- 答え
abc363a r = 100 - mod r 100

B - Japanese Cursed Doll

「呪われた日本人形」は Cursed Japanese Doll
Japanese Cursed Doll は「日本の呪い人形」つまり藁人形のこと（うそ）

問題 ABC363B

シグネチャを決める。横着する。

abc363b :: [Int] -- N,T,P
        -> [Int] -- Li
        -> Int   -- 答え

愚直に

データの範囲が狭いので、力業で計算する。
$D$ 日後の髪の長さは $L_i + D \geq T$ つまり $L_i \geq T - D$ を満たす $i$ の個数を毎回数え、それが $P$ 以上になる最小の $D$ を線形探索する。

abc363b [_n,t,p] ls = head
  [ d
  | d <- [0 ..], let td = t - d
  , p <= length (filter (td <=) ls)]

計算機に優しく

$i$ の長さが $T$（以上）になる日は $D_i = \max(0, T - L_i)$ なので、そのような$D_i$に1を置いた配列の累積和をとり、$P$以上になる位置を探す。（累積和による積分）

import Data.Array

abc363b :: [Int] -> [Int] -> Int
abc363b [_n,t,p] ls =
  length $ takeWhile (p >) $      -- p以上になる日数を数える
  scanl1 (+) $ elems $            -- 累積和をとって
  accumArray (+) 0 (0, t)         -- 投票をカウント
  [(max 0 $ t - l, 1) | l <- ls]  -- iがTになる日に1投票

計算量は $O(N+T)$

もっと入力範囲が大きいときは、配列でなく IntMap Int を使って飛び飛びな投票を集計する。

公式解説のやり方

また、この問題はそれぞれの人が何日後に髪の長さが $T$ 以上になるか記録し、ソートを用いて小さい方から $P$ 番目の値を求めることで、$O(N \log N)$ の計算量で解くことができます。

$D_i = \max(0, T - L_i)$ をソートして $P$ 番目の値を求める。

import Data.List

abc363b :: [Int] -> [Int] -> Int
abc363b [_n,t,p] ls = max 0 $ sort [t - l | l <- ls] !! pred p

あっすごくシンプル。
「$P$番めの値」はソートして $O(N \log N)$ としなくても、クイックセレクトで $O(N)$ でできることになっているけど。

C - Avoid K Palindrome 2

問題 ABC363C

シグネチャを決める。横着する。

abc363c :: Int    -- N
        -> Int    -- K
        -> String -- S
        -> Int    -- 答え

たかだか $N \leq 10$ なので、$10! = 3628800$ 個を Data.List.permutations で数え上げ、重複を除き、長さ $K$ の回文を含むものを除き、残ったものを数える愚直でいけるだろう、とナメてかかったら

import Data.List
import qualified Data.Set as S

abc363c :: Int -> Int -> String -> Int
abc363c n k s = length . filter prop . S.toList . S.fromList . permutations $ s
  where
    prop t = not . any paliK . take (n - pred k) . tails $ t
    paliK x = reverse (take k x) `isPrefixOf` x

まるで間に合わなかった。
Data.List.permutations は、辞書順も重複も気にせず謎の順序で生成するから使いにくい。

公式解説では案の定 C++ の next_permutation で、文字の重複について配慮して列挙して数えればよいとあるのでそうすると

import Data.List
import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad

abc363c :: Int -> Int -> String -> IO Int
abc363c n k s =
  do
    mv <- MUV.new n
    forM_ (zip [0..] $ sort s) (uncurry (MUV.write mv))
    loop mv 0
  where
    jmax = pred $ div k 2
    loop mv cnt = do
      c <- prop mv
      let !cnt1 = if c then succ cnt else cnt
      b <- MUV.nextPermutation mv
      if not b then return cnt1 else loop mv cnt1
    prop mv =
      and <$> forM [0 .. n - k] (\i ->
        or <$> forM [0 .. jmax] (\j -> do
          a <- MUV.read mv (i+j)
          b <- MUV.read mv (i+pred k-j)
          return $ a /= b
          )
        )

TLEが2つ残ってしまった。

真面目に生成する

文字の重複を許す順列を全て生成するには、「未使用で残っている文字」全てで総当たりする代わりに、「未使用で残っている文字種」だけに関して総当たりする。
なので、どの文字がいくつ残っているかを把握しておく。

全て生成してから任意の位置に長さKの回文を含むものを除く代わりに、生成の途中で長さKの回文が完成した瞬間にそこより深い再帰を全て枝刈りする。

import Data.List
import Data.Array

abc363c :: Int -> Int -> String -> Int
abc363c n k s = recur n "" $ accumArray (+) 0 ('a','z') [(c,1) | c <- s]
  where
    recur :: Int            -- 生成残り文字数
          -> String         -- 生成済みの文字列（逆順）
          -> Array Char Int -- 文字の使える残り個数
          -> Int            -- 答え
    recur 0 _ _   = 1
    recur m t arr = sum
      [ recur (pred m) t1 arr1
      | (c,d) <- assocs arr, d > 0
      , let t1 = c : t
      , n - pred m < k || notHeadPali t1
      , let arr1 = arr // [(c, pred d)]
      ]

    -- 前K文字が回文でない
    notHeadPali t1 = not (reverse (take k t1) `isPrefixOf ` t1)

文字を仮定しない

文字を仮定すると配列で数えることができるが、そうでない場合でも何とかなる。
要素の Ord 順で、重複を意識した辞書順生成と、回文フィルタとで分離した形にしてみる。
つまり枝刈りなし。

import Data.List

abc363c :: Int -> Int -> String -> Int
abc363c n k s = length $ filter prop $ myPermutations s
  where
    k2 = div k 2
    prop xs = not $ any headPali $ take (n - pred k) $ tails xs
    headPali xs = reverse (take k2 xs) `isPrefixOf` drop (k - k2) xs

myPermutations :: (Ord a) => [a] -> [[a]]
myPermutations = loop . group . sort
  where
    loop [] = [[]]
    loop xss = [x : res | (x, xss1) <- rot xss, res <- loop xss1]
    rot (xs:xss) = (head xs, cons (tail xs) xss) : [(y, xs:yss) | (y,yss) <- rot xss]
    rot [] = []
    cons [] xss = xss
    cons xs xss = xs:xss

C問題に七転八倒してしまった。

D - Palindromic Number

ここまできて、363が回文であることに掛けているのだと気づいた。

問題 ABC363D

シグネチャを決める。
$10^{18}$ は Int64 ぎりぎりの数なので、そんな $N$ 番目の回文数は Integer でないと表せない。しかし数として扱う必要もないので String で結果を返すことにする。

abc363d :: Int    -- N
        -> String -- 答え

1桁の回文数は 0～9 の10個
2桁の回文数は 11～99 の9個
3桁の回文数は A=1～9 0～9 A の 90個
4桁の回文数は A=1～9 B=0～9 B A の 90個
5桁の回文数は A=1～9 B=0～9 0～9 B A の 900個

以下同様に考えて、

• 偶数 $2k$ 桁の回文数は $9 \times 10^{k-1}$ 個あり、
  その $i$ 個め（0始まり）は $10^{k-1} + i$ とその reverse を連結した数
• 奇数 $2k+1$ ($k \geq 1$)桁の回文数は $90 \times 10^{k-1}$ 個あり、
  その $i$ 個めは $10 \times 10^{k-1} + i$ とその reverse を連結した数

とわかる。そして、
$m$ 桁以上の $i$ 個めの回文数とは、
$m$ 桁の回文数の個数 $C_m$ 未満 $i < C_m$ ならその $i$ 個めで、
さもなくば $m+1$ 桁以上の回文数 の $i - C_m$ 個めである。

結果

abc363d :: Int -> String
abc363d n
  | n <= 10   = show (pred n)
  | otherwise = evenMode 1 (n - 11)

evenMode, oddMode :: Int -> Int -> String

-- evenMode k i : 偶数桁 2k で、0始まりi個めの回文数
evenMode k i
  | i >= ub   = oddMode k (i - ub)
  | otherwise = s ++ reverse s
  where
    base = 10 ^ pred k
    ub = 9 * base
    s = show $ base + i

-- oddMode k i : 奇数桁 2k+1 で、0始まりi個めの回文数
oddMode k i
  | i >= ub   = evenMode (succ k) (i - ub)
  | otherwise = s ++ tail (reverse s)
  where
    base = 10 * 10 ^ pred k
    ub = 9 * base
    s = show $ base + i

E - Sinking Land

問題 ABC363E

シグネチャを決める。横着する。

abc363e :: [Int]   -- H,W,Y
        -> [[Int]] -- Aij
        -> [Int]   -- 答え

二次元配列の書き換えが強く要求されそうなこういう問題は気後れする。

immutableに試し実装

状態として

• $G$ ground 水没していない、地上マスの一覧 (座標の Set）
• $S$ seaside 海に面した、水没しうるマスの一覧
  • 標高の低い順に取り出したいので、優先度付きキューがいい？
  • キューに重複して登録されて問題無いようにしたいので、標高と座標の対の Set で表す
• $T$ t 現在の時刻（上昇した海面の高さ）

を管理する。
時刻が $Y$ に至るまで、以下の処理を繰り返すことで状況をシミュレートできる。

• キュー $S$ が空のとき、この先沈むマスはない。残りの時間だけ $|G|$ を繰り返し出力して終了する。
• キュー $S$ の先頭のマスの座標を $p$ その標高を $A_p$ とする。
  • $T < A_p$ なら、まだしばらく沈まないので、そうなる時刻まで今の値 $|G|$ を繰り返し出力し、時計を $A_p$ に進める。
  • $T \geq A_p$ なら、水没モードに入る。
    $p$ は水没したので、$G$ から除く。
    $p$ の4近傍のマス $q$ で、水没していないもの $q \in G$ かつ、水没しうるもの $A_q \leq Y$ を $S$ に追加する。
    この $q$ や、あるいは無関係なキューの後続で 水没するものを引き続き水没させるため、この処理に再突入する。

import Data.List
import Data.Array
import qualified Data.Set as S

abc363e :: [Int] -> [[Int]] -> [Int]
abc363e [h,w,y] ass = loop 1 seaside0 ground0
  where
    bnds = ((1,1),(h,w))
    a = listArray bnds $ concat ass
    seaside0 = S.fromList $
      [(a ! p, p) | i <- [1..h], let p = (i,1), a ! p <= y] ++
      [(a ! p, p) | i <- [1..h], let p = (i,w), a ! p <= y] ++
      [(a ! p, p) | j <- [1..w], let p = (1,j), a ! p <= y] ++
      [(a ! p, p) | j <- [1..w], let p = (h,j), a ! p <= y]
    ground0 = S.fromDistinctAscList $ range ((1,1),(h,w))
    loop :: Int -- 今から検討する時刻=水面高さ
         -> S.Set (Int, (Int,Int)) -- 海に面した陸地の高さ順、座標
         -> S.Set (Int, Int) -- 陸地マスの集合
         -> [Int] -- 答え、各時刻の陸地マス数
    loop t seaside ground
      | S.null seaside = replicate (y - pred t) (S.size ground)
      | t < ap         = replicate (ap - t) (S.size ground) ++ loop ap seaside ground
      | otherwise      = loop t seaside1 (S.delete p ground)
      where
        ((ap, p@(i,j)), ss1) = S.deleteFindMin seaside
        seaside1 = foldl' (flip S.insert) ss1
          [ (a ! neighbor, neighbor)
          | neighbor <- [(pred i, j), (i, pred j), (i, succ j), (succ i, j)]
          , S.member neighbor ground
          , a ! neighbor <= y ]

WAにはならなかったので、結果は正しいようだが、もちろんTLEする。

mutable arrayでACを取りに行く

IOモナドで命令型配列を使う。
結果は得られた側から自分で print する。

ground :: Set (Int,Int) を IOUArray (Int,Int) Bool に置き換える。
また、not でしか参照しなかったので、論理を反転させて undersea に改名する。

seaside もフラグ配列にし、それとは別にキューも保持する。これで先頭のアクセスが $O(1)$ になる。

キューにしたので重複登録しないようにする。（というかハマった。）
immutable実装の seaside0 は角を二重に登録しているうえに、$H=1$や$W=1$といった異常な状況では猛烈に二重登録をしてしまう。

結果：1728ms, 223MB

ユーザ解説(1) Union-Find で解く

陸地に加えて周囲1マスの海も考え、マスをノードとするメッシュ状のグラフを考える。
ただし、初期状態では海の頂点だけが存在し、互いに接続されているとする。

次から、時刻 $T$ に、標高 $T$ なマスを表す頂点をグラフに追加し、隣接するマスがあるならそれとつなぐ辺も追加する。
これが、海に面していれば海の分割とunionされ、内陸のくぼ地なら別の分割となる。
なので、海の分割のノード数から、水没したマスの個数がわかり、陸地のマスの数もわかる。

さらに水面が上昇してダムが決壊すると、くぼ地の分割が海とunionされる、という仕掛け。

なるほど、こういうUnion-Findの使い方もあるのね。

ユーザ解説(2) ダイクストラ法で解く

ダイクストラ法では普通は辺の重みを足し合わせることで最短距離を求めるが、ここでは $\max$ を取ることで、辺の重みの最大値が最も小さくなる経路のその最大値を求めることができる、というアイデア。（まずこの段階で宇宙猫になってしまう。）

上のUnion-Findのようなメッシュ状のグラフを考え、辺の重みを両端のマスの標高の大きい方とする。
これで、海 $(0,0)$ からの各マスへの経路で、経由するマスの標高を最も低くしたときの最大値がわかる。その値はつまり、その経路を通ってそのマスに水が流入するときの海面高さに他ならない、と。

なんと、こういうダイクストラ法の応用もあるのね。

そしてさらっと書かれる

本質的には公式解説と等価です。

の一言。殺傷力高い。

自分の解法もキューを用いているので、こちらの方法に類型分類されるのかな、と。

F - Palindromic Expression

問題 ABC363F

シグネチャを決める。横着する。

abc363f :: Int     -- N
        -> String  -- 答え

Eと比べてこういう問題は好み。

数 $X$ について、

• それが数字0を含まない回文数なら、そのままそれが答え（のひとつ）
• そうでないとき、その0を含まない約数 $b$ ($X = bX_1$ とする) と $b$ を逆にした（123を321にするという意味）数 $d$ が $X_1$ の約数 ($X_1 = dX_2$ とする) という $b,d$ について、$X_2$ が答え $T$ を持つなら、$b \times T \times d$ が答え

と再帰的に求められる。
$N \leq 10^{12}$ と大きいように見えるが、気にするべき $X$ は $N$ の約数に限られ、$b, d$ の値はさらにそこから条件を満たすものだけなので、見た目より軽い。

結果

import qualified Data.IntMap as IM

abc363f :: Int -> String
abc363f n =
  case im IM.! n of
    (s:_) -> s
    []    -> "-1"
  where
    facn = factors n
    im = IM.fromDistinctAscList [(k, f k) | k <- facn] -- f のメモ化
    facn0 = [ (k,s,kr)                                  -- b,dの候補
            | k <- takeWhile ((n >=) . (^ 2)) $ tail facn -- 約数1を除外し、b≦dの範囲で
            , let s = show k, notElem '0' s
            , let kr = read (reverse s), mod n (k * kr) == 0 ]
    f x
      | reverse s == s, notElem '0' s = [s]
      | otherwise = take 1 cands -- ひとつで十分ですよ
      where
        s = show x
        cands = [ t ++ '*' : fx2 ++ '*' : reverse t
                | (k,t,kr) <- facn0
                , (x1,0) <- [divMod x k]
                , (x2,0) <- [divMod x1 kr]
                , fx2 <- im IM.! x2 ]

-- @gotoki_no_joe
-- 約数列挙
factors :: Int -> [Int]
factors 1 = [1]
factors n = 1 : loop 2 [n]
  where
    loop k us
      | k2 >  n =     us
      | k2 == n = k : us
      | r  == 0 = k : next (q:us)
      | True    =     next    us
      where
        (q,r) = divMod n k
        next = loop (succ k)
        k2 = k * k

G - Dynamic Scheduling

実行時間制限 8 sec！TLE答案を提出するとひたすら待つハメになる。

解説がやたらと生えて、頭の痛くなる専門用語が飛び交っているので、火傷しないうちに撤収。