ABC339をHaskellで

お久しぶりです。

A - TLD

問題 ABC339A

シグネチャを決める。

abc339a :: String  -- S
        -> String  -- 答え

結果

Preludeにあるリスト関数では、ひっくり返してから takeWhile するのが楽そう。

abc339a = reverse . takeWhile ('.' /=) . reverse

Data.List.elemIndices を使えば、.が出現する最後の位置を取り出せる。

import Data.List

abc339a :: String -> String
abc339a s = drop (succ $ last $ elemIndices '.' s) s

命令型の場合は、文字列を後ろから見て.が見つかったらそれより後ろだけを結果とするのが前者、最後に.を見かけた位置を覚えつつ文字列全体をスキャンして、最後に見た位置より後ろを結果とするのが後者で、割と考え方も対応している。

B - Langton's Takahashi

問題 ABC339B

シグネチャを決める。

abc339b :: Int -- H
        -> Int -- W
        -> Int -- N
        -> [String] -- 答え

ラングトンの蟻がぼっちでトーラス空間に閉じ込められている。

結果

純粋計算にこだわり、「黒く塗られたマスの座標の集合」を更新するという方針でやってみる。

import qualified Data.Set as S

abc339b :: Int -> Int -> Int -> [String]
abc339b h w n =
  [ [if S.member (i,j) stN then '#' else '.' | j <- [0 .. pred w]]
  | i <- [0 .. pred h]]
  where
    stN = loop n (0,0) (-1,0) S.empty

    -- n 残りステップ数カウントダウン
    -- ij 現在位置(0始まり)
    -- d 現在の向き
    -- fld 黒座標集合
    loop 0 _ _ fld = fld
    loop n ij d fld = loop (pred n) ij1 d1 fld1
      where
        black = S.member ij fld
        fld1 = (if black then S.delete else S.insert) ij fld
        d1 = turn black d
        ij1 = add ij d1

    add (a,b) (c,d) = (mod (a+c) h, mod (b+d) w)

    turn True  (i,j) = (-j,i)
    turn False (i,j) = (j,-i)

旋回はこのように、回転行列 $\left ( \begin{array}{ll}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array} \right )$ の $\theta = \pm 90^\circ$ でするのが好き。

というテクニックも便利。公式解説もこれを使っている。
真似をするなら、turnは捨てて、loopの末尾2行を次に差し替える感じか。

        d1 = mod ((if black then pred else succ) d) 4
        ij1 = add ij ([(-1,0),(0,1),(1,0),(0,-1)] !! d1)

そもそも問題の内容が命令型の配列を指示とおりに更新する話なので、Data.Array.ST の練習問題と思って書いてもよいだろう。

C - Perfect Bus

問題 ABC339C

シグネチャを決める。

abc339c :: Int   -- N
        -> [Int] -- Ai
        -> Int   -- 答え

人数の折れ線グラフを考えて、一番低い箇所がちょうど無人となるように上下にずらしたときの最後の値が答え。

結果

abc339c _n as = last ss - minimum ss
  where
    ss = scanl (+) 0 as

フレンズさんの別解

ユーザ解説 by kyopro_friends

もしシミュレーションの過程でバスに乗っている人数が負になった場合、”実はバスにはもっと人が乗っていた”と思うことで、その時点でバスに乗っている人数は0だったことにします。

では、$A_i$の走査は一度だけで済む。

abc339c _n = foldl (\s a -> max 0 (s + a)) 0

D - Synchronized Players

問題 ABC339D

シグネチャを決める。

abc339d :: Int      -- N
        -> [String] -- Si
        -> Int      -- 答え

状態として二人のプレイヤーの座標を考えると、その状態数は $N^4 \leq 60^4 = 12,960,000$ と結構大きいので、純粋なデータ構造で管理するのはつらい。素直に mutable array で調査済みの状態を管理して、幅優先探索をする。

結果

import Data.Array
import Data.Array.ST

import Control.Monad
import Control.Monad.ST

type Point = (Int,Int)
type P2 = (Point,Point)

-- 単純な1マス移動
mn, me, mw, ms :: Point -> Point
mn (i,j) = (pred i, j)
me (i,j) = (i, succ j)
mw (i,j) = (i, pred j)
ms (i,j) = (succ i, j)

abc339d :: Int -> [String] -> Int
abc339d n ls = runST action
  where
-- isOK : 移動可能なマスかの判定
    bnds = ((1,1),(n,n))
    fld = listArray bnds [c /= '#' | l <- ls, c <- l]
    isOK ij = inRange bnds ij && fld ! ij
-- fで移動可能なら移動、さもなくば留まる
    move f ij = let ij1 = f ij in if isOK ij1 then ij1 else ij
-- 二人の初期位置
    [ij00, ij01] = [(i,j) | (i,l) <- zip [1..] ls, (j,'P') <- zip [1..] l]
-- 本体
    action :: ST s Int
    action = do
      flgs <- newArray (((1,1),(1,1)),((n,n),(n,n))) False
      loop flgs 0 [(ij00,ij01)] []
-- 幅優先探索
    loop :: STUArray s P2 Bool -- 調査済みの状態フラグ
         -> Int                -- 歩数
         -> [P2]               -- 調査キュー
         -> [P2]               -- 次の周回での調査キュー
         -> ST s Int           -- 答え
    loop _ _ [] [] = return $ - 1 -- 万策尽きた
    loop flgs steps [] news = loop flgs (succ steps) news [] -- 次の周回へ
    loop flgs steps ((ij,kl):p2s) news = do
      flg <- readArray flgs (ij, kl)
      case () of -- multiway if
        _ | flg       -> loop flgs steps p2s news -- 調査済みならばスルー
          | ij == kl  -> return steps             -- 重なっていたら成功
          | otherwise -> do
              writeArray flgs (ij,kl) True
              writeArray flgs (kl,ij) True  -- 面倒なので両方ともフラグを立てる
              loop flgs steps p2s news1
      where
        news1 =
          [ (ij1, kl1)
          | f <- [mn,me,mw,ms]
          , let ij1 = move f ij
          , let kl1 = move f kl
          , (ij,kl) /= (ij1, kl1)] ++ news

E - Smooth Subsequence

問題 ABC339E

シグネチャを決める。

abc339e :: Int    -- N
        -> Int    -- D
        -> [Int]  -- Ai
        -> Int    -- 答え

最長上昇部分列(LIS)を求めるDPと同じ流れで、部分列の末尾の値をキー、その値を末尾として作れる、条件を満たす最長の部分列の長さを値とする対応付けを、前から順に更新することを考える。
次の値$A_i$に注目したとき、$[A_i - D, A_i + D]$の範囲のキーの値の最大値（値が一つもないときは0とする）+1を$A_i$に対応づける。
全ての$A_i$について見終わったら、対応付けに含まれる値の最大値が答えとなる。

しかし、上の「範囲のキーの値の最大値をとる」処理が$O(ND)$となり、特に$D$が大きいとき、素朴なマップや配列を用いた実装では間に合わない。

そしてこれは判りやすいセグメント木の出番なのでそうする。

結果

未定義な箇所は minBound で表した。実際には-1で用は足りる。

セグメント木の実装は省略するので全体は提出を参照。

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

abc339e :: Int -> Int -> [Int] -> Int
abc339e n d as = runST action
  where
    action = do
      st <- makeSegTree max minBound 500001 (minBound :: Int)
      foldM_ (\_ ai -> do
        x <- querySegTree st (ai - d) (succ ai + d)
        setSegTree st ai (succ $ max 0 x)
        ) () as
      querySegTree st 0 500001

data SegmentTree s a = SegmentTree Int (a->a->a) a (STree s a)

type STree s a = MUV.MVector s a

-- makeSegTree f u n i :: 0からn-1の要素を初期値iで初期化したセグメント木を作る
makeSegTree :: MUV.Unbox a => (a->a->a) -> a -> Int -> a -> ST s (SegmentTree s a)

-- setSegTree st i x :: iをxに書き換える
setSegTree :: MUV.Unbox a => SegmentTree s a -> Int -> a -> ST s ()

-- querySegTree st l r :: [l,r)区間の問い合わせ
querySegTree :: MUV.Unbox a => SegmentTree s a -> Int -> Int -> ST s a

自分用備忘録：
セグメント木は2分木で、葉の枚数は$N$、根の方に登る毎にノード数は半分になるので全ノード数は$2N$、木の高さは$\log_2 N$
最初、ノード数を$N^2/2$と勘違いして、$A_i \leq 5 \times 10^5$のセグメント木は無理、と思い込んでしまった。（それは$N \times N$の三角行列）

雑感

公式解説の

$dp_{i,j}​$を 各$(i,j)$について陽に持つのではなく配列の使い回しをする

フレンズさんのtweet

in-placeに更新

どちらも、$i$と$j$二つの添え字を使って説明されたDPは、二次元配列に置くことを強い原則として考えているらしい言い方なのがいつも気になっていて、添え字$i$の値を作るのに添え字$0～i$全ての値を使う一般の場合はそうだけど、この問題は最初から、直前の値だけが必要で、終わったら捨てられることは明らかなので、そこはいちいち強調する必要あるのかしらと。

F - Product Equality

問題 ABC339F

シグネチャを決める。

abc339f :: Int        -- N
        -> [Integer]  -- Ai
        -> Int        -- 答え

$A_i < 10^{1000}$と、一千桁の数が1000個も渡される。
特に何か、数字列として扱うだけで何とかなる秘策がある訳ではなく、多倍長整数も無いような言語はここでさようなら、ということらしい。
確率的な計算だけでやる、という条件付きで、多倍長整数なしでできるというのが出題意図だったのかも、ということで追記します。

お試し

TLE覚悟で試してみる。せめて、同じ数の重複分についてはあらかじめ数えて、あとは総当たりでする。

import qualified Data.Map as M

abc339f :: Int -> [Integer] -> Int
abc339f _n as = sum
    [ ci * cj * ck
    | (ai, ci) <- M.assocs m
    , (aj, cj) <- M.assocs m
    , Just ck <- [M.lookup (ai * aj) m]
    ]
  where
    m = M.fromListWith (+) [(a, 1) | a <- as]

繰り返しの回数を半分にすることもしないこの実装が、1511ms, 20MBで通ったのでこれでお終いでもいいのだけど。

高速化

高速化のポイントがふたつある。

• $A_j$ を走査する内側のループは、$A_i$を走査する外側のループの値以上の範囲だけを調べれば済む。
  ただし、$A_i \neq A_j$ のとき、$A_i > A_j$ の側の場合の数を数えるために結果を倍にする必要がある。
• $A_i$の最大値を$A_x$として、$A_j$の走査は$A_x / A_i$を超えたら打ち切ってよい。

import qualified Data.Map as M
import Data.List

abc339f :: Int -> [Integer] -> Int
abc339f _n as = sum
    [ ci * cj * ck
    | ((ai, ci):ajcjs) <- tails $ M.assocs m  -- AjはAiより後ろだけ
    , let ub = div amax ai
    , (aj, cj) <- takeWhile ((ub >=) . fst) $ -- 上限で打ち切り
                  ((ai, ci) :) $              -- Aj=Aiも試す
                  map (fmap (2 *)) ajcjs      -- Aj≠Aiは結果を倍にする仕込み
    , Just ck <- [M.lookup (ai * aj) m]
    ]
  where
    (amax,_) = M.findMax m
    m = M.fromListWith (+) [(a, 1) | a <- as]

これで111msに高速化された。

出題者の意図

モジュロ演算で $A_i \times A_j = A_k \Rightarrow (A_i \bmod x) \times (A_j \bmod x) = (A_k \bmod x)$ から、条件に合わないことはCPU整数の範囲で高速に判定できる。
$x$を複数切り替えて試すとき、十分な確率で、モジュロ演算のみで、条件を満たすことも判定できる、ということらしい。

こういう確率的な判定を使うときは、高速で確率的な判定を通ったものについて、厳密な検査で確定させるという手順が続くと思っていたのだが、公式解説では20個程度の素数$x$を試せばよいとし、フレンズさん解説に至ってはたった2個の値でためすだけで、厳密な検査なしで済ませている。

さらにひっかかるのが、公式解説の、これが「正当な解法である」「xを素数にしなくとも正解できる」というくだり。これはあくまで「競技プログラミングのコンテストで固定的テストケースを用いた判定において、不正解と見做されずに済む」だけで、「いかなる入力に対しても正しい結果を返す、真に正しいプログラム」ではないと思うのだけど。

追記：出題者の意図２

多倍長整数として読み込んで剰余を計算する、ではなく、数字列を整数に読み込む段階で剰余をとりつつやれば、CPU整数だけでもできることについて。

複数の割る数を用意して、それらで割った余りの組をキーとしてマップに個数を数えて、後は同じ。

ただ、確率的なプログラムなのがなぁ。

import Data.Char
import qualified Data.Map as M

divs :: [Int]
divs = [ 2^31 - 1
       , 2^31 - 19 -- prime
       ]

-- 数字列で受け取るのでシグネチャ変更
abc339f :: Int -> [String] -> Int
abc339f _n as = sum
    [ ci * cj * ck
    | (rsi, ci) <- M.assocs m
    , (rsj, cj) <- M.assocs m
    , let rijs = [mod (ri * rj) d | (ri,rj,d) <- zip3 rsi rsj divs]
    , Just ck <- [M.lookup rijs m]
    ]
  where
    m = M.fromListWith (+) [([readMod m a | m <- divs], 1) | a <- as]

-- 剰余を取りながら数字列を読み込む
readMod :: Int -> String -> Int
readMod m ds = foldl step 0 ds
  where
    step acc d = mod (acc * 10 + digitToInt d) m

除数はこちらから拝借しました。
というかこのJavaコードがまさにそういう解法になっていた。

除数をいくつにするかというマシュマロ

G - Smaller Sum

問題 ABC339G

シグネチャを決める。

abc339g :: Int     -- N
        -> [Int]   -- Ai
        -> Int     -- Q
        -> [[Int]] -- αi,βi,γi
        -> Int     -- 答えBi
abc339g n as q abcs = ...

（ふんわりとした方針は思いついたもので合ってたが、間に合う実装は解説を見ないと書けなかった。）

考える

セグメント木ライクな

とある区間に関する何かを求めたいとなると、セグメント木が連想される。（今回２度目）
しかし単なる区間の和ではなく、$X_i$以下の値に限った値が欲しい。

一つの数列に固定して、様々な$X_i$以下の値の総和を求めたいという部分問題を考える。
数列の値$A_i$の頻度表をイメージすると、$X_i$以下の部分の面積（に値をそれぞれ乗じたもの）を高速に得るには、$A_i$をキーにした累積和の表を作っておけばよい。

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

solve :: [Int] -- Ai
      -> [Int] -- Xi
      -> [Int] -- 答え {Ai}のXi以下の要素の総和
solve as xs = [maybe 0 snd $ IM.lookupLE x im | x <- xs]
  where
-- リスト指向
    as1 = group $ sort as
    im = IM.fromDistinctAscList $ zip (map head as1) $ scanl1 (+) $ map sum as1
-- IntMap重点
    im = snd $ IM.mapAccum step 0 $ IM.fromListWith (+) [(a,a) | a <- as]
    step a c = let x = a + c in (x, x) 

この表 im を、セグメント木の全ての区間に対して貼り付けておき、クエリに対して、最小限の区間の組み合わせで答えを求める。
セグメント木では、ノードに貼り付けてある情報を統合することで答えを見つけるが、この問題では、一つのクエリ$X_i$について注目しているときに、全ての値に対応できる累積和表IntMapを構築する必要はないので、木を降りていき、それぞれのIntMapに問い合わせた結果を足し合わせる。

累積和表の統合

個々の区間に対する累積和表を、毎回 as を直接参照して作るのはおそらく無駄が多い。これをやると$N$要素を$\log N$回舐めることになる。
二つの子のノードに付いた、もしくはそれを作るのに使った情報を利用して、親ノードを作ることを考える。

リスト指向の場合、group する前の、区間の要素を整列したリストを残すと、それを（いわゆるマージソートの〉マージで、$O(N)$で統合できる。ソートし直しで$O(N \log N)$かけるより速いだろう。このアイデアが、公式解説の "merge-sort tree" に相当するものと思われる。

merge :: Ord a => [a] -> [a] -> [a]
merge xxs@(x:xs) yys@(y:ys) =
  case compare x y of
    LT -> x : merge xs yys
    GT -> y : merge xxs ys
    EQ -> x : y : merge xs ys
merge xs [] = xs
merge [] ys = ys

IntMap指向の場合も、assocs で取り出したリストを統合できる。どちらのマップも累積和を持っているので、直前の相手側の値、つまりこれまでの最大値を補えばよい。

merge :: Int -> [(Int,Int)] -> Int -> [(Int,Int)] -> [(Int,Int)]
-- xlv, ylvは最後に見た値、初期値は0
merge xlv xcxcs@((x,c):xcs) ylv ydyds@((y,d):yds) =
  case compare x y of
    LT -> (x, ylv + c) : merge c xcs ylv ydyds
    GT -> (y, xlv + d) : merge xlv xcxcs d yds
    EQ -> (x, c + d)   : merge c xcs     d yds
merge xlv [] _ylv yds = map (fmap (xlv +)) yds
merge _xlv xcs ylv [] = map (fmap (ylv +)) xcs

ハイブリッド構造

区間の長さが1になるまで、セグメント木を完全に構築すると、累積和表IntMapが$2N$個作られる。これはいささか多すぎる。（実際TLE, MLEになってしまう。）
クイックソートを、要素数が少なくなったところでバブルソートに切り替えるのと同様に、ここでも、区間がそれなりに短くなったところで、直接 as から答えを求めるモードに切り替えることにする。
設定する値は、大きいと端の計算が重くなるし、小さいとIntMapを作りすぎるので、程ほどに。

chunkLen :: Int
chunkLen = 128

中途半端なMerge-Sort Treeの構築

以上の方針で、途中で打ち止めるmerge-sort treeを構築する。
（マージソートはどこにも入っていないが、名前はそのままにしておく。）

配列を木にするのでなく、代数的データ型で木を作った。

data MergeSortTree = MST Int (Array Int Int) MST -- 要素数、Aiの配列、木

data MST = Leaf                          -- 打ち切りノード。直接数える
         | Node (IM.IntMap Int) MST MST  -- 累積和表、左右の部分木

makeMST :: Int -> [Int] -> MergeSortTree
makeMST len xs = MST w xarr $ fst $ iter 0 w w
  where
-- 要素数以上の2の冪
    w = until (len <=) (2 *) 1
-- Aiを保存する配列
    xarr = listArray (0, pred w) $ xs ++ repeat 0
-- [p,q)区間の幅wに対するMSTノードを再帰的に作る、累積和表も添える
    iter :: Int -> Int -> Int -> (MST, IM.IntMap Int)
    iter p q w
      | w <= chunkLen = (Leaf, im0)
      | otherwise     = (Node im lt rt, im)
      where
        im0 = snd $ IM.mapAccum step 0 $ IM.fromListWith (+) $
              [(a,a) | i <- [p .. pred q], let a = xarr ! i]
        w2 = div w 2
        m = p + w2
        (lt, lm) = iter p m w2
        (rt, rm) = iter m q w2
        im = IM.fromDistinctAscList $ merge 0 (IM.assocs lm) 0 (IM.assocs rm)
    step a c = let x = a + c in (x, x) 

クエリに応える

再帰で降りていくやり方自体はセグメント木と同じ。

queryMST :: MergeSortTree -> Int -> Int -> Int -> Int -- [左,右)で範囲を指定
queryMST (MST w xarr mst) a b x = loop 0 w w mst
  where
    -- loop p q w t : 現在位置 i の左右 [p,q) その幅 w
    loop p q _w Leaf = sum -- Leafなら直接計算、範囲外なら自然にiが空
      [ ai
      | i <- [max a p .. pred $ min b q]
      , let ai = xarr ! i, ai <= x]
    loop p q w  (Node im l r)
      | q <= a || b <= p = 0 -- 領域外
      | a <= p, q <= b   = maybe 0 snd $ IM.lookupLE x im -- 完全に包含されていれば表を引く
      | otherwise        = loop p m w2 l + loop m q w2 r
      where
        w2 = div w 2
        m = p + w2

起動部

全体を駆動するところで、前の結果を使って次のクエリを暗号解除する。

abc339g :: Int -> [Int] -> Int -> [[Int]] -> [Int]
abc339g n as _q abcs = tail $ scanl anstep 0 abcs
  where
    t = makeMST n as 
    anstep b0 abc = queryMST t (pred l) r x
      where
        l:r:x:_ = map (xor b0) abc

セグメント木の添え字が0始まりなので、$L_i, R_i$は l, succ r をずらして pred l, r になる。
最終結果は2626ms, 391MB

純粋関数型計算でGが書けたのはうれしい。

平方分割？

ユーザ解説 by potato167とユーザ解説 by kyopro_friendsが「平方分割」を説明していて、後者の最後に、公式解説に繋がると書いてある。

でももう疲れたので、また気が向いたら考えるということで、今日はここまで。