ABC350をHaskellで

ここのところ後半が難しくて詰まっていたのですが、今回は急に肩透かしくらいましたね。
G問題の別解を追記しました。

A - Past ABCs

問題 ABC350A

シグネチャを決める。

abc350a :: String -- S
        -> Bool   -- 答え

6文字中後ろ3文字が数字で、それが、316でない 1～349 であるかを確認する。

結果

abc350a s = 0 < n && n < 350 && n /= 316
  where
    n = read $ drop 3 s

B - Dentist Aoki

問題 ABC350B

シグネチャを決める。

abc350b :: Int   -- N
        -> Int   -- Q
        -> [Int] -- Ti
        -> Int   -- 答え

頼りになる僕らの味方 accumArray の出番。

結果

import Data.Array

abc350b :: Int -> Int -> [Int] -> Int
abc350b n _q ts = length $ filter id $ elems $
                  accumArray f True (1,n) [(t, ()) | t <- ts]
  where
    f b _ = not b

C - Sort

問題 ABC350C

シグネチャを決める。

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

添え字iを1から順に、違う場所にある値iを位置iに持ってくるような交換をしていく。
運良く位置iに値iが既に来ていたら、それはスキップできる。
最大$N-1$回の交換を行うことになる。

位置iにある値は何か、という普通の向きの情報と、
値iが現状の列でどの位置にあるか、という逆引き情報とを、
同期させて管理する。

結果

やることは明確なので、速度のために初めから mutable array で書く。

import Data.Array.ST
import Control.Monad
import Control.Monad.ST

abc350c :: Int -> [Int] -> [[Int]]
abc350c n as = runST $
  do
    i2a <- newListArray (1,n) as :: ST s (STUArray s Int Int) -- 位置iからそこに今ある値a
    a2i <- newArray_ (1,n)       :: ST s (STUArray s Int Int) -- 値aが今ある位置i
    zipWithM_ (writeArray a2i) as [1..]
    loop i2a a2i 1
  where
    loop _ _ i | n < i = return []
    loop i2a a2i i = do
      ai <- readArray i2a i
      if ai == i then loop i2a a2i (succ i) else do
        j <- readArray a2i i -- 位置iに入れるべき値が今ある場所
--        writeArray i2a i i -- 位置iにiを入れる、この処理は省略できる
--        writeArray a2i i i
        writeArray i2a j ai  -- iのあった位置に、先頭にあった値を移す
        writeArray a2i ai j
        ([i,j] :) <$> loop i2a a2i (succ i)

公式の「Pythonで解く場合の注意点 by mymelochan」は公式解説に書いてあることと何も違わない、と思ったが、Pythonの配列に index() という、この問題で使えそうに見える落とし穴メソッドがあるよって話なのね。

ユーザ解法にある別解

「 A以外の配列を作成しない方法 by hirayuu_At」について。
公式のアプローチは、前から見ていって現在の位置 $i$ に正しい値 $i$ を今ある位置 $j$ から持ってくる、そのために今位置 $i$ にある正しくない値 $A_i$ をそこに捨てる、という向き。
これを逆に、現在の位置 $i$ にある迷子の値 $j$ を、正しい位置 $j$ に帰してあげる、そのために今位置 $j$ にある値 $A_j$ は位置 $i$ で引き取る、という向き。下手すると現在の位置が進まなそうな気がするけれど、$j$ の方はどんどん揃っていくので、スワップの回数は同じで済む。

abc350c :: Int -> [Int] -> [[Int]]
abc350c n as = runST $
  do
    arr <- newListArray (1,n) as :: ST s (STUArray s Int Int)
    loop arr 1
  where
    loop _ i | n < i = return []
    loop arr i = do
      j <- readArray arr i
      if j == i then loop arr (succ i) else do
        aj <- readArray arr j  -- jをしまうべき場所に今ある値
        writeArray arr j j
        writeArray arr i aj
        ([i,j] :) <$> loop arr i

D - New Friends

問題 ABC350D

シグネチャを決める。$A_i, B_i$ は横着する。

abc350d :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- Ai, Bi
        -> Int      -- 答え

問題文で「$(A_i, B_i)$ は相異なる」だけだと、逆向きの辺で意地悪できるかもと思ったけれど、前の行に $A_i < B_i$ があったのでその心配は不要だった。

グラフの連結成分を求めて、それぞれの連結成分について、ノード数から、完全グラフが持つ辺の数が求められる。
また、既に張ってある辺の本数は数えられる。その差が、追加できる辺の本数。

結果

これまでだと、自作Union-Findが登場する流れだけど、今まで存在は把握していたけれどスルーしていた Data.Graph を使ってみる。
両方向に辺を張ることで無向グラフを作ると、強連結成分、をお任せで計算できる。
連結成分は薔薇木 Tree で返されるのだが、その要素の取り出しが見当たらなかったので手書きした。

import Data.Graph

abc350d :: Int -> Int -> [[Int]] -> Int
abc350d n m abs = subtract m $ flip div 2 $ sum $ map f $ scc g
  where
    g = buildG (1,n) [e | a:b:_ <- abs, e <- [(a,b),(b,a)]] -- 双方向(無向)グラフ
    f t = k * pred k                            -- ノード数kの完全グラフの辺の数×2
      where
        k = length vs $ nodes t

nodes :: Tree a -> [a] -- 木のノードリスト
nodes t = loop t []
  where
    loop (Node a ts) rest = a : foldr loop rest ts

E - Toward 0

問題 ABC350E

シグネチャを決める。全体的に横着する。

abc350e :: [Int]   -- N,A,X,Y
        -> Double  -- 答え

$a \div b = \lfloor a / b \rfloor$ とする。
$N$ の値に対して、そこから0にするまでのコストの期待値 $f(N)$ を考える。
$f(0) = 0$ である。
正の $N$ について、$f(N)$ は $A$ を使う方だとして $f_A$とおく。$X$ 円払うので $f_A = X + f(N \div A)$ である。
サイコロを使う方だとして、その期待値を$f(N) = f_B$ とする。1が出たときに $\lfloor N/1 \rfloor = N$ で進まない恐れがある。
$f_B = Y + \frac{1}{6} \sum_{b=1}^6 f(N \div b)$
$6f_B = 6Y + f_B + \sum_{b=2}^6 f(N \div b)$
$f_B = \big \{ 6Y + \sum_{b=2}^6 f(N \div b) \big \} / 5$
実際にはこの小さい方を選択するべきである。$f(N) = \min(f_A,f_B)$

$1 \leq N \leq 10^{18}$ と巨大なので、$f$ を配列に持つDPをストレートに実行できない。
ここで実際には、$N$ は $1/2$ ～ $1/6$ とかなり高速に減少するので、飛び飛びの要素にしか用事はなく、メモ化再帰で計算すれば求められる。

結果

Stateモナドに既知の対応をメモしたIntMapを持ち回らせる。

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

abc350e :: [Int] -> Double
abc350e [n,a,x,y] = evalState (f n) (IM.singleton 0 0)
  where
    dx = fromIntegral x
    dy6 = 6 * fromIntegral y
    f k = do
      im <- get
      case IM.lookup k im of
        Just v -> return v
        Nothing -> do
          cb <- (/ 5) . (dy6 +) . sum <$> mapM (f . div k) [2 .. 6]
          ca <- (dx +) <$> f (div k a)
          let c = min ca cb
          modify (IM.insert k c)
          return c

似たような問題を割と最近に見たような。

公式解説から

調べるべき $N$ の値は、$m = 2^p 3^q 5^r$ と書ける整数 $m$ の $N \div m$ だけでよいと。
（その理由が $(N \div a) \div b = N \div (ab)$ とあるが、なんかズレることないのかと引っかかる。）
これがわかれば、攻めのDPで実装できる。

abc350e :: [Int] -> Double
abc350e [n,a,x,y] = im IM.! n
  where
    dx = fromIntegral x
    dy6 = 6 * fromIntegral y
    im = IM.fromList $
         (0, 0) :
         [ (n3, min ca cb)
         | n1 <- takeWhile (0 <) $ iterate (flip div 2) n
         , n2 <- takeWhile (0 <) $ iterate (flip div 3) n1
         , n3 <- takeWhile (0 <) $ iterate (flip div 5) n2
         , let ca = dx + im IM.! div n3 a
         , let cb = (dy6 + sum [im IM.! div n3 b | b <- [2 .. 6]]) / 5
         ]

F - Transpose

問題 ABC350F

シグネチャを決める。

import qualified Data.ByteString.Char8 as BS

abc350f :: BS.ByteString  -- S
        -> IO ()          -- 答えはその都度出力

考える

定義通りに実行していたら時間が足りるはずもなく。
カッコの中について、大文字と小文字を入れ替えるとともに、順序も逆にする。
二度この処理を行うと、元の列に戻る。

(abc(def)ghi(jkl)mno)
->
(abc FED ghi LKJ mno)
->
 ONM jkl IHG def CBA

defの区間とjklの区間は、2度処理したため元に戻っている。ただし周囲は1度だけ処理しているのため、反転している他、defとjklの出現する位置も元とは異なる。

abcのような、カッコのない区間のブロックごとに考えて、カッコの入れ子がなす木構造をデータ型に取り出す段階と、木構造を降りて、前から順に出力を実行する段階とに分けて考える。

abc350f = output . buildTree

data Tree = Node [Tree] | Leaf BS.ByteString

buildTree :: BS.ByteString -> Tree
output :: Tree -> IO ()

読み込み

Attoparsecを使ってみよう。

import Data.Char
import Data.Attoparsec.ByteString.Char8

まず、アルファベットの並び1文字以上のByteString区間を切り出せたら、Leafにできる。

    leaf = Leaf <$> takeWhile1 isAlpha

これだと多分、読み込み済みのByteStringの一部を指すポインタとして文字列が得られて速い。

カッコ開きが見つかったら、カッコ閉じまで、0個以上の木を読み取り、Nodeが作れる。

    node = do
      char '('
      ts <- many' tree
      char ')'
      return $ Node ts

treeは、leafとnodeのどちらでもよい。

    tree = choice [node, leaf]

トップレベルは、treeの1つ以上の繰り返しで、それもNodeに包んで返すことにする。

    toplevel = Node <$> many1 tree

attoparsecは「受理成功」「失敗」の他に、「構文解析が途中だから入力の続きを供給してほしい」という結果をとりうる。
many系パーサは、手持ちの入力が尽きても、続きが供給されたらまだ続く可能性があるので、topLevelパーサは終わってくれない。
つまりparse toplevel bsはDone i rにならずPartial ...になってしまう。

マニュアルには parseOnly (myParser <* endOfInput) とやれば止まるとあるのでそうする。
（でも、manyが貪欲な場合、endOfInputが効くところまで抜けてこないような気がするんだが、よくわからない。）

buildTree :: BS.ByteString -> Tree
buildTree bs = Node res
  where
    Right res = parseOnly (many1' tree <* endOfInput) bs
    ...

書き出し

Treeデータ型のNodeを降りるたびにモードを切り替えて、前から無加工で出力するか、後ろから変換しつつ出力する。

output :: Tree -> IO ()
output = recur True
  where
    recur True  (Node vs) = mapM_ (recur False) vs            -- 前から順に出力
    recur False (Node vs) = mapM_ (recur True) $ reverse vs   -- 逆転して出力
    recur False (Leaf bs) = BS.putStr bs                      -- 前から順に出力
    recur True  (Leaf bs) = BS.putStr $ BS.reverse $ BS.map inv bs -- 逆転して出力

inv :: Char -> Char
inv c
  | isUpper c = toLower c
  | otherwise = toUpper c

前から出力するモードのNode区間について、Leafにはモードが既に逆転して与えられるため、スイッチがFalseのときがLeafを前から出力するべきタイミングになる。

buildTreeをparsecなしに手書きすることもできるが、多分最初からこれでやった方が楽だろう。

公式解説のやり方

の説明は何ともややこしいが、ユーザ解説 by ngtkanaの説明

最も左から、右を向いてスタートして

• 基本は向いている方向に進む
• (, ) を踏むたびに対応する括弧にジャンプして、向いている方向を逆にする
• 向いている方向が左のときだけ lowercase/uppercase を反転

の説明は明解ですばらしい。
indexを進めていく、その動作は完全に命令型な計算なので、気が向いたら。⇒できました。

もう一つのユーザ解説「平衡二分探索木を使うことで簡単に解くこともできます。」はどういう方針なのかさっぱりわからない。

G - Mediator

問題 ABC350G

何やらメモリの制限が厳しく設定されている。時間制限の方は少し緩い。

考える

メモリの節約のため、クエリを一つずつ読み込んでは処理するスタイルでいこう。
foldMの状態変数は、直前のタイプ2の結果$X_k + 1$の値。

main :: IO ()
main = do
  [n,q] <- bsGetLnInts
-- 前準備
  foldM_ (\xk1 _ -> do
    [a,b,c] <- map succ . zipWith (flip mod) [2,n,n] .
               map ((flip mod 998244353) . (xk1 *)) <$>
               bsGetLnInts
    case a of
      1 -> {- タイプ1の処理 -} >> return xk1
      2 -> -- タイプ2の処理
    ) 1 [1 .. q]

何が出題の狙いかわからないので、素朴なやり方を考えてみる。
頂点ごとに、それが接続している頂点の集合を配列で維持する。
タイプ1では $u$ に $v$, $v$ に $u$ を追加する。
タイプ2では、$u$ の隣接頂点と $v$ の隣接頂点の共通部分を探す。これはたかだか一つしかない。

ノードが $N \leq 10^5$ クエリが $Q \leq 10^5$ で、メモリの必要量が $10^{10}$ オーダーになるように見えるが、二次元配列でなく、効率的な IntSet のような形式で持つなら、全体で $Q$ で抑えられるから間に合いそう。

ということで素直にやってみる。

前準備

1から$N$の頂点について、隣接頂点番号の IntSet を持つ mutable array を作る。

-- インポート
import Data.Array.IO
import qualified Data.IntSet as IS

-- mainの中
  g <- newArray (1,n) IS.empty :: IO (IOArray Int IS.IntSet)

タイプ1

上述のように登録する。

-- mainのcaseの中
      1 -> mode1 g b c >> return xk1

mode1 :: IOArray Int IS.IntSet -> Int -> Int -> IO ()
mode1 g b c = do
  gb <- readArray g b
  writeArray g b (IS.insert c gb)
  gc <- readArray g c
  writeArray g c (IS.insert b gc)

タイプ2

素直に解を探し、出力する。またその解+1を返す。

-- mainのcaseの中
      2 -> mode2 g b c

mode2 :: IOArray Int IS.IntSet -> Int -> Int -> IO Int
mode2 g b c = do
  gb <- readArray g b
  gc <- readArray g c
  let ans = head $ (++ [0]) $ IS.elems $ IS.intersection gb gc
  print ans
  return $ succ ans

結果

381ms, 33MBでAC。
何だったのだろう。

フレンズさんの解説

X

アライグマ「G問題はクエリ平方分割なのだ！
フェネック「Link-Cut Treeっていうすごいデータ構造を使うとクエリを直接処理することもできるみたいだねー」

そうでもなかったけど。

追記：ユーザ解説の方法

マージテク O(N log N) 解法 by toamについて。

Union-Findは木の根も辺もどんどん書き換えるけれど、似て非なる、指定された辺を忠実に張り、ただし根を勝手に決めて、各頂点に関して親へのリンクを管理する。これを par[] とする。

これが維持されていれば、タイプ2の解は、

• u → w → v となっているときの w （par[par[u]] = v のとき par[u]）
• v → w → u となっているときの w （par[par[v]] = u のとき par[v]）
• u → w ← v となっているときの w （par[u] = par[v] のとき par[u]）

のいずれかで得られる。

このpar[]をタイプ1について維持するために、指定された u,v に関して、u → v とした上でuの属する木の全てのノードをvの方に向け直すか、逆に v → u でするか、どちらかをする必要がある。

ここで、解説では、

• タイプ1で向きを選ぶために、u,v の属する木のサイズが小さい方を書き換える側として選ぶ。
  それを判断するために、Union-Findも同時に維持する。
• 小さい方のノード全てを、深さ優先探索で回ってpar[]を更新する。
  それを行うために、既知の全ての辺を頂点に紐付けて記録する。

としている。

自分の改善

木のどこかにある（一般に根ではない）uの par[u] = v と向け直して、同様にuの属する木のpar[]を全てvの方に向け直した結果がどうなるかを考えてみると、

• uより葉のノードについては、変化がない。
• uから根までの経路上のノードは、逆向きになる。
• 根から、uとは違う方に生えている全ての枝は、変化がない。

ので、結局、uから根までを更新すれば済み、深さ優先探索とか全ての辺の記録は不要で、par[]さえ有れば実行できる。

また、そうなると、根からuまでの距離と、根からvまでの距離を比べて、短い方を更新することが理にかなっているとわかる。木全体のサイズは不要で、Union-Findも使わずにできる。

という方針で書き直してみたが、どうもTLEする。考えた結果、

• 根からuまでの距離を数える
• 根からvまでの距離を数える
• 両者のどちらが小さいかを比較する

という手順でやると、片方が異様に高い木なとき、その高さを数え終わるまで計算が終わらない。
どちらがより低いか、だけが知りたいことなので、par[]をたぐる動作をuからとvからを同時に実行し、どちらかが根に到着した時点で終了できる。高い方の木の高さを具体的に数える必要はなかった。

この改善を追加することで、補助的な情報を全く記録しないで動く版ができた。

• C++, 23ms, 3.7MB （←toamさんのオリジナル版は 55ms, 10MB）
• Haskell, 51ms, 14MB
• Python, 340ms, 89MB

おわりに

Data.Graph や Data.Attoparsec に手を出すことができてよかった。
今回はやり過ごせてしまった難しそうなキーワード「クエリ平方分割」「Link-Cut Tree」に警戒しておく必要があるか。