ABC344をHaskellで

A - Spoiler

問題 ABC344A

シグネチャを決める。

abc344a :: String -- S
        -> String -- S

結果

-- 手書きの再帰で
abc344a s = loop1 s
  where
    loop1 ('|':s) = loop2 s
    loop1 ( c :s) = c : loop1 s
    loop2 ('|':s) = s
    loop2 ( _ :s) = loop2 s

-- Preludeで
abc344a s = a ++ c
  where
    (a, _:b) = span ('|' /=) s
    _:c = dropWhile ('|' /=) s

B - Delimiter

問題 ABC344B

数値列が入力されるが、長さ$N$があらかじめ与えられる代わりに、0終端を検出して逆順で出力せよという要求。これまでのAtCoderの出題傾向と少し毛色が違って、入出力まで含めたコンピュータの操作としてのプログラミングを求めている。それは、Haskellで書くのが楽しくない部分。そもそもシグネチャを決められない。

結果

なのでこちらも、mainの再突入とかでひねくれてみる。

import Control.Monad

main :: IO ()
main = do
  a <- getLine
  when (a /= "0") main
  putStrLn a

C - A+B+C

問題 ABC344C

シグネチャを決める。
引数が妙に多いので、どうせリストの長さは使わないので外す。

abc344c :: [Int] -- Ai
        -> [Int] -- Bi
        -> [Int] -- Ci
        -> [Int] -- Xi
        -> [Bool] -- 答え

$N \times M \times L$通りの組み合わせでの和を先に求めておき、$X_i$が含まれるかを調べる。
全ての組み合わせを作るのも、力任せにやっても間に合うけれど。

結果

import qualified Data.IntSet as IS

abc344c :: [Int] -> [Int] -> [Int] -> [Int] -> [Bool]
abc344c as bs cs xs = [IS.member x is | x <- xs]
  where
-- 力任せ
    is = IS.fromList [a + b + c | a <- as, b <- bs, c <- cs]
-- 初めからIntSetで
    is = mul cs $ mul bs $ IS.fromList as
    mul xs is = IS.unions [IS.map (x +) is | x <- xs]

D - String Bags

問題 ABC344D

シグネチャを決める。
どうせリストの長さは使わないので$A_i$は外す。

abc344d :: String     -- T
        -> Int        -- N
        -> [[String]] -- Sij
        -> Int        -- 答え

$len(S)$を文字列$S$の長さとする。
$cost[i][j]$ を、袋1から$i$まで使って、$T$の前$j$文字を作成する、最もコストの低いやり方のコストとしてDPする。
$cost[0][0] = 0$ で、$cost[N][len(T)]$が求める答えとなる。
$T$の$j$文字め以降が$S_{i,k}$と一致するとき、これを使って文字列の作成を進めることができるので、$cost[i][j + len(S_{i,k})]$ はそのような $cost[i][j] + 1$ の最小値である。また、袋$i$を使わずに済ます場合も含めて $cost[i][j]$ の候補に $cost[i-1][j]$ も入る。

結果

IntMapを使うことで、当てはまる値かない場合を含めて扱い、（珍しく）配るDPで実装した。

import qualified Data.IntMap as IM
import Data.List

abc344d :: String -> Int -> [[String]] -> Int
abc344d t n sss = IM.findWithDefault (-1) (length t) imZ
  where
    imZ = foldl step (IM.singleton 0 0) sss
    step im ss = foldr ($) im
        [ IM.insertWith min (k + length s) (succ c)
        | (k, c) <- IM.assocs im
        , let t1 = drop k t
        , s <- ss
        , isPrefixOf s t1
        ]

E - Insert or Erase

問題 ABC344E

シグネチャを決める。

abc344e :: Int      -- N
        -> [Int]    -- Ai
        -> Int      -- Q
        -> [[Int]]  -- query_i
        -> [Int]    -- 答え

考える（少しおふざけ）

クエリ1は、既存の$x$の「次」に新規の値$y$を挿入する。
同じ$x$に対して何度も挿入があると、直前に挿入したら$y$との間にさらに次の値がどんどん挟まっていく。
現状の数列の様子を何らかのデータ構造に保存する必要があるが、クエリで与えられる$x$がそのデータ構造の中で（今）どこにあるかを見つけるための逆写像的なインデックスも同時に維持する必要があると考えられる。

数列の様子を入れるためのデータ構造として、辞書のようなものが使えないか考えてみる。辞書式順序で、二つの語$x,z$の間に入る語$y$は、末尾の文字を変えるだけでなく、文字を追加しても不可能になってしまうことがある。例えば aaaとaaaaの間に入る語はない。

実数はみっちり詰まっているので、$x,z$の位置をそれぞれ実数$s,t$で表すと、$(s+t)/2$は必ず両者の間になる。これをキーとしてMapを作ればよいように思うが、Doubleは有限精度でとても間に合わない。

しかし我々にはData.Ratio.Rationalという無限精度の分数があった！これをつかってみよう。

import qualified Data.Map as M
import qualified Data.IntMap as IM
import Data.Ratio

abc344e :: Int -> [Int] -> Int -> [[Int]] -> [Int]
abc344e _n as _q qus = init $ M.elems $ fst $ foldl step (p2x0, x2p0) qus
  where
-- 数直線上にAiを配置
-- Aiは(i-1)%1とし、A_Nの次にダミー0を置いておく
    p2x0 =  M.fromDistinctAscList
      [(p, a) | (p,a) <- zip [0 % 1 ..] $ as ++ [0]] :: M.Map Rational Int
-- Aiを置いた位置Rationalを持つ逆写像
    x2p0 = IM.fromList [(a, p) | (p,a) <- zip [0 % 1 ..] as]
-- クエリ2では両方からxを除く
    step (p2x, x2p) (2:x:_) = (M.delete (x2p IM.! x) p2x, IM.delete x x2p)
-- クエリ1では次の要素との中間点にyを置く
    step (p2x, x2p) (1:x:y:_) = (M.insert py y p2x, IM.insert y py x2p)
      where
        px = x2p IM.! x
        Just (pz,_) = M.lookupGT px p2x
        py = (px + pz) * (1 % 2)

結果 TLEx6
同じ位置にひたすら挿入されると、Rationalの分母が$2^Q$になるのでコーナーケースでは辛いが、大抵の状況では間に合ってしまいそうな。

真面目にやる

上の、逆写像も維持するのはそのまま、列の様子を保存するデータ構造として、IntMapのような木（immutableにやさしい）ではなく、双方向リンクリストを使えというのが出題意図だと読み取れる。
この問題も問題Bと同様に、命令型バリバリな設定が今までと違う雰囲気を感じた。

書き換え可能な配列を使って、

• 添え字は0から$N+Q$まで
  • 0はリンクリストの先頭と末尾を指す特別な要素
  • $i = [1,N]$は$A_i$を順に入れる
  • $i = [N+1, N+Q]$は、クエリ$j$が挿入のとき$N+j$に$y$を入れる
  • つまりガベージコレクションとかしない
• 要素は3項組で
  • 手前の要素の添え字
  • 自分の値
  • 次の要素の添え字

という双方向リンクリストと、逆引き表をクエリに応じて更新する。

Data.Vector.Unboxed.Mutableを使った結果 927ms, 147MB
でもHaskellでこんなことしても楽しくないにゃあ。
Rustも双方向リンクリスト苦手なんだっけ？

F - Earn to Advance

問題 ABC344F

シグネチャを決める。

abc344f :: Int      -- N
        -> [[Int]]  -- Pij
        -> [[Int]]  -- Rij
        -> [[Int]]  -- Dij
        -> Int      -- 答え

それぞれのマスに到達するやり方で、(1)なるべく早い、(2)なるべく所持金が多い、(3)通過した経路の$P_{i,j}$の最大値が最も大きい、ようなものを選ぶDPをすればいい、という大枠は見当が付いたが、枝刈りをする基準がわからなくて解説を見た。

自分のイメージでは、到達時刻をX軸、所持金をY軸に点をとり、そこから、$P_{i,j}$の最大値の傾きで伸びる半直線を描き、これが露出している区間が検討するべき場合に相当すると考えた。これはユーザ解説 by MMNMMと近いように思う。ただし思っただけで実装していない。

結局、マスの位置$(i,j)$と、経路上の$P_{i,j}$の最大値を添え字にして、到着時刻の昇順、次に所持金の降順で最適なもの一つを保持するDPでよいようだ。

公式解説 by kyopro_friendsでは、マスの位置$(i,j)$については配列で、$P$の最大値はIntMapのキーにして、そのバリエーションは$P_{i,j}$の種類が上限なので要素数は$N^4$となる、とやっている。
ユーザ解説 by evimaは、$P$の最大値をキーにする代わりに、その最大値の$P_{i,j}$のあるマスの座標$(i,j)$を配列の添え字として、直接$N^4$要素の配列を張っていることに相当するのだと思う。

集めるDPで暗黙のDP

arrがDP配列で、マスの位置$i,j$を添え字に、内容は$P$の最大値から、到着時刻と所持金の対へのIntMapである。

import Data.Array
import qualified Data.IntMap as IM

abc344f :: Int -> [[Int]] -> [[Int]] -> [[Int]] -> Int
abc344f n pss rss dss = minimum $ map fst $ IM.elems $ arr ! (n,n)
  where
    n1 = pred n
    bnds = ((1,1),(n,n))
    p = listArray bnds $ concat pss            -- Pij
    r = listArray ((1,1),(n,n1)) $ concat rss  -- Rij
    d = listArray ((1,1),(n1,n)) $ concat dss  -- Dij
    arr = listArray bnds $ map f $ range bnds  -- DP配列
    f (1,1) = IM.singleton (p ! (1,1)) (0,0)   -- 出発点
    f ij@(1,_) = right ij
    f ij@(_,1) = down  ij
    f ij       = IM.unionWith cmp (right ij) (down ij)
    cmp tm1@(t1,m1) tm2@(t2,m2) = -- (time,money) の良い方を選びだす
      case compare t1 t2 of
        LT -> tm1
        GT -> tm2
        _  -> if m1 >= m2 then tm1 else tm2
    right (i,j) = walk (i,j) r (i, pred j)
    down  (i,j) = walk (i,j) d (pred i, j)
    walk ij w ij0 = IM.fromListWith cmp -- ij0から向きwの移動コストを払ってijに至る場合
      [ (max pij pmax, (succ t + k, m - wij + k * pmax))
      | (pmax, (t, m)) <- IM.assocs $ arr ! ij0 -- ij0の全ての場合について
      , let k = divrup (max 0 $ wij - m) pmax]  -- 最も稼げるpmaxで不足分を稼ぐのにk日かかる
      where
        pij = p ! ij  -- ijでの稼ぎPij
        wij = w ! ij0 -- ij0からの移動コスト

-- 切り上げ除算
divrup x y = negate $ div (negate x) y

見通しはよくて、「らしい」コードになったが、TLEx2になってしまう。

MVectorを使って普通に配る命令型のDPで実装したら通った。
$P$の値を昇順に背番号をつけて、この番号をDP配列の3つめの添え字にしている。

G - Points and Comparison

問題 ABC344G

シグネチャを決める。

abc344g :: Int      -- N
        -> [[Int]]  -- Xi, Yi
        -> Int      -- Q
        -> [Int]    -- G0, Ra, Rb
        -> Int      -- 答え
abc344g n xys q [g0,ra,rb] = ...

もうなんもわからんので早々に公式解説を見た。

公式解説の考え方

• とある傾き$A_k$で切片$B_k$の直線は、切片つまりY軸との交点は$B_k$
• それぞれの点$(X_i,Y_i)$を通り傾き$A_k$の直線の、Y軸との交点がそれより上にあるものの数を数える

ということで、各 $(X_i, Y_i)$ に対して切片 $Y_i - A_iX$ を全て求めておくことができれば、二分探索で$B_k$以上のものの個数は数えられる、のは予想していた。

• $A_k$を変化させていったとき、切片そのものは連動して変化してしまうけれど、各点の切片の順序関係はそれほど変化しない
• 二つの点$(X_i,Y_i)$と$(X_j,Y_j)$を考えたとき、その傾き$\frac{Y_i-Y_j}{X_i-X_j}$の前後で、両者の切片の並び順が変化する
• なので切片の順序を傾きの変化に追従させて、その並び順について二分探索して個数を数える

これは思い至らなかった。ただ、同じ位置関係に3つ以上の点が並んだりすると厄介そうだなと。
そしてさらに解説では、

• 初めは、点列を辞書順に並べておく（これは傾きが$-\infty$のときの切片の順序に相当する）
• 監視する傾きは、この並びで隣接する点だけでよい、また、X座標が等しい点どうしは、切片の順序の入れ替えは起きないので気にしない
  • これは、傾きに応じて切片が線形に変化したとき、先頭どうしが最初に衝突するから、当たり前といえば当たり前、蟻の行列どうしが衝突するのは先頭の蟻どうしに決まっている
• 入れ替えが発生したら、前にずれた点のそのさらに前にいる点との位置関係、後ろにずれた点のそのさらに次にいる点との位置関係が、次に変化する傾きに達したときに入れ替えを起こせるようにイベントを追加する必要がある
  • 先頭の蟻どうしかすれ違ったら、次に起きるのはその次の蟻とすれ違うこと、という話

そこで、傾きの大きさを有理数で優先度付きキューのキーとして、監視する傾きに達したときに入れ替えるべき2点（切片の順序で隣接しているはず）を特定できる情報を値として持っておく。
クエリについても、$A_k$の昇順に直しておく。
今から$A_k, B_k$について調べるとき、前準備として、$A_k$以下の傾きに関するキューの要素を全て処理する。処理により波紋が生じるので、それはキューに追加する。

こんな大きな絵は見えなった。
ただ、実装を読んでみると、キューにとっておいたのにもかかわらず、いざ交換しようとしたら見失っていたときはスルーする、みたいな変な動作が混じっていて、微妙に怪しい。
3つ以上の点が同じ角度で並んだときは、結局、全体の順序が一気に入れ替わる（一直線状に並んだ電柱を、電柱から右に首を出して見たときと、左から覗いたときで、左右の並びが逆転する）ので、そうなるように交換の順序を保つ（安定でないソートを使うとここが破綻する）ことも、かなり薄氷を踏んでいる感じ。

自分の考え

点群は増えたりしないので、全ての点どうしの組み合わせたかだか$N^2 \ (N \leq 5000)$個について、あらかじめ、入れ替える傾きを全て計算してしまえばよいのでは。
そうすれば途中で要素を追加したりしないので、優先度付きキューもいらないのでは。

切片の昇順での点列を、配列 Array Int (Int,Int) で保持する。
初期の辞書順での順位を背番号として、背番号から、現在の並びでどの位置にいるかをひく表 Array Int Int を連動して維持する。

    xya = listArray (1,n) $ sort [(x,y) | x:y:_ <- xys]
    i2k = listArray (1,n) [1 .. n]

全ての（X座標が異なる）2点間の傾きをキーに、それで入れ替わる点の番号の対を値としてキューを作る。

    queueL :: [(Ratio Int, Int, Int)]
    queueL = sort $ (10^9 % 1, 0, 0) : -- sentinel
      [ ((yj - yi) % (xj - xi), i, j)
      | i <- [1 .. pred n], let (xi,yi) = xya ! i
      , j <- [succ i .. n], let (xj,yj) = xya ! j, xi < xj]

クエリも辞書順に整列して、$A_i,B_i$のタプルで表す。

    queueR :: [(Ratio Int, Int, Int)]
    queueR = [(a % 1, a, b) | (a,b) <- abs]
-- 乱数列Gi
    gs = tail $ iterate (\g -> mod (48271 * g) modBase) g0
-- クエリ(Ai,Bi)をQ個
    abs = sort $ take q $ map abfun $ iterate (drop 3) gs
    abfun (g2:g1:g0:_) = ( mod g2 (ra + succ ra) - ra
                         , mod (g1 * modBase + g0) (rb + succ rb) - rb)
    modBase :: Int
    modBase = 2147483647 -- 2^31 - 1

以上の仕込みを元に、クエリのキューqueueRを順に処理する。
そのとき、必要な分の点交換のキューqueueLを先に処理する。
queueLの末尾には番兵を立たせたので尽きる心配はない。

    ans = loop xya i2k queueL queueR 0
    loop _ _ _ [] acc = acc
    loop xya0 i2k0 qL0@((aL,i,j):qL1) qR0@((aR,a,b):qR1) acc
      | aL <= aR  = loop xya1 i2k1 qL1 qR0 acc  -- swap
      | otherwise = loop xya0 i2k0 qL0 qR1 acc1 -- count
      where
-- 点iと点jの位置kを交換する
        ki = i2k0 ! i
        kj = i2k0 ! j
        i2k1 = i2k0 // [(i, kj), (j, ki)]
        xya1 = xya0 // [(ki, xya0 ! kj),(kj, xya0 ! ki)]
-- b以上になる点の個数を数えて足し込む
        acc1 = acc + n - x
        (x,_) = binarySearch prop 0 (succ n)
        prop k = y - a * x >= b
          where
            (x,y) = xya0 ! k

配列がimmutable arrayなので実行速度は望むべくもないが、例を入れると正しい結果が返る。
なので、このアルゴリズムをSTArrayやVector.Unboxed.Mutableにしてみたが、時間制約、空間制約を満たすことができなかった。Vector.Algorithmsとか使ってないのでよくわからない。
後は「速い」Haskellを書くのが得意な方にお任せします。

このままでは引き下がれないので、アルゴリズムをそのままつたないC++に翻訳した。
結果：3576ms, 746MB

雑感

他の提出を見てみると、プラグマをいじって最適化を強化する必要はあるものの、問題文のとおりに計算をするだけのアプローチで、10秒制限をギリギリ突破できるらしい。
コード長がめちゃめちゃ短いので目立つ。
速い言語だとそういうこともできてしまう一方で、Javaですら現時点でACが0とは、この出題はちょっと問題があったのでは。

要素を入れ替えて二分探索とか、効率的にimmutableに書けない操作が要求されてもにょり、今回はB問題、E問題も異質で、いろいろ違和感でした。