ABC308 A～G をHaskellで

多重集合がやたらと出てきた、ような。

A - New Scheme

問題 ABC308A

シグネチャを決める。

abc308a :: [Int]  -- S1～S8
        -> Bool   -- 答え

• リストの隣接要素について何かするときはzipWithとtailを使うイディオムがある。
• 条件1で昇順に並んでいるなら、下限と上限を全ての要素について調べる必要はない。

結果

abc308a xs = cond1 && cond2 && cond3
  where
    cond1 = and $ zipWith (<=) xs $ tail xs
    cond2 = 100 <= head xs && lst xs <= 675
    cond3 = all ((0 ==) . flip mod 25) xs

B - Default Price

問題 ABC308B

シグネチャを決める。

abc308b :: Int       -- N
        -> Int       -- M
        -> [String]  -- Ci
        -> [String]  -- Di
        -> [Int]     -- Pi
        -> Int       -- 答え

結果

Data.Mapを使う。

import qualified Data.Map as M

abc308b :: Int -> Int -> [String] -> [String] -> [Int] -> Int
abc308b _n _m cs ds (p0:ps) = sum [M.findWithDefault p0 c dp | c <- cs]
  where
    dp = M.fromList $ zip ds ps

C - Standings

問題 ABC308C

シグネチャを決める。

abc308c :: Int           -- N
        -> [(Int, Int)]  -- Ai, Bi
        -> [Int]         -- 答え

コイントスの成功率 $\frac{A_i}{A_i + B_i}$ の降順、等しいものどうしは番号の昇順で整列するということだが、率の方を浮動小数点数で扱うと誤差で間違えるようなテストケースがあるのだろう。
丁寧に通分して、$\frac{A}{B}$ と $\frac{C}{D}$ の比較を $\frac{AD}{BD}$ と $\frac{BC}{BD}$ で行う（しかも降順に）比較関数をソートに渡すことを想定しているように見える。

結果

Haskellには有理数があるのでそれを使う。ただし率をマイナスにして降順で整列させる。

import Data.Ratio

abc308c :: Int -> [(Int,Int)] -> [Int]
abc308c _n abs = map snd $ sort [(- a % (a + b), i) | ((a,b), i) <- zip abs [1..]]

Cにしては入力データが大きくて、Stringでのんびり読み込みしていたら1900msでギリギリだった。
ちゃんとByteStringですれば778msだったので、気をつけよう。

計算精度の話

誤差評価 by kyopro_friends によると、C++ の long double なる拡張倍精度浮動小数点数の精度ならこの出題の誤差を上回るらしい。
残念ながらNumeric.LongDoubleはAtCoderにないし、そもそもそういう手段を使うべきでもない。
また、「10^20を分子に掛けてから整数で割り算するだけで精度が出る」とも書いてあるが、これは整数がデフォルトで多倍長なPythonならではの発想か。

D - Snuke Maze

問題 ABC308D

シグネチャを決める。

abc308d :: Int       -- H
        -> Int       -- W
        -> [String]  -- Si
        -> Bool      -- 答え

次に進めるマスに書いてある文字を追跡しながら、幅優先探索でゴールまで行けばいいだけ。
開始位置の文字が s でない場合すらあるので微妙に注意。

結果

いつものように二重ループを一つの再帰関数で。
到達済みのマスを記録するのに Set (Int,Int) は重いので、ij2x で番号を振って IntSet を使う。

import Data.Array
import qualified Data.IntSet as IS

abc308d :: Int -> Int -> [String] -> Bool
abc308d h w ss = bfs (cycle "snuke") IS.empty [(1,1)] []
  where
    m = listArray (0, h * w - 1) $ concat ss -- 迷路も一次元配列で持つ
    bfs :: String      -- 今踏むべき文字が先頭なsnukeの無限の繰り返し
        -> IS.IntSet   -- 到達済みのマスの番号
        -> [(Int,Int)] -- 探索するべき座標のリスト
        -> [(Int,Int)] -- 次のステップで探索するべき座標のリスト
        -> Bool        -- 答え、ゴールに到達できるか
    bfs _ _ [] [] = False -- 次のステップに進みたいが空なのは行き詰まり
    bfs (_:cs) visited [] kls = bfs cs visited kls [] -- 次があるなら文字も進めて進行
    bfs cs@(c:_) visited ((i,j):ijs) kls -- 個々の座標を処理
      | IS.member x visited = bfs cs visited ijs kls -- 到達済みなら捨てる
      | m ! x /= c          = bfs cs visited ijs kls -- 文字が違うなら捨てる
      | i == h, j == w      = True                   -- ゴールなら成功
      | otherwise           = bfs cs visited1 ijs (kls1 ++ kls) -- 4近傍を次のステップに追加
      where
        x = pred i * w + pred j
        visited1 = IS.insert x visited
        kls1 = [(pred i,j) | 1 < i] ++ [(succ i,j) | i < h] ++
               [(i,pred j) | 1 < j] ++ [(i,succ j) | j < w]

E - MEX

問題 ABC308E

シグネチャを決める。
$S$が長いのでByteStringを使う。

import qualified Data.ByteString.Char8 as BS

abc308e :: Int            -- N
        -> [Int]          -- Ai
        -> BS.ByteString  -- S
        -> Int            -- 答え

考える

$A_i$ は0から2の3通りの場合しかなく、$\text{mex}(A_i,A_j,A_k)$ の値は27通りの場合に対して計算しておける。

import Data.Array

pa = listArray ((0,0,0),(2,2,2))
     [ head $ [x | x <- [0..3], notElem x [ai, aj, ak]]
     | ai <- [0..2], aj <- [0..2], ak <- [0..2]
     ]

3つの要素からなる話は、真ん中を基準に考える。
文字Eがあるようなとある位置 $j$ について、
そこより手前で文字Mがある位置について、0,1,2がそれぞれいくつあるか、
そこより後で文字Xがある位置について、0,1,2がそれぞれいくつあるか、
がわかれば、paの要素にその個数を掛けたものが位置 $j$ のEに関して獲得するポイントの全てである。
つまり、現在の位置の数字 $A_j$ をajとし、Mの位置におげる0,1,2の個数を順に持つリストm012とXに関するリストx012から、

sum [ ci * ck * pa ! (ai,aj,ak)
    | (ci,ai) <- zip m012 [0..]
    , (ck,ak) <- zip x012 [0..]
    ]

として算出できる。
これを、Sを走査して順に求めていけば全体のポイントを計算できる。

• 現在の位置の文字が M のときは m012 の A_i の位置を1増やす。
• 現在の位置の文字が E のときは、上のポイントを計上する。
• 現在の位置の文字が X のときは、x012 の A_k の位置を1減らす。

m012は初期値オール0から始めて、Mに遭遇するたびに増やしていけばよい。
x012の初期値はS全体を調べる必要があり、Xに遭遇するたびに減らしていくと最後にオール0になるものである。

結果

128ms, 47MBでACした。

F - Vouchers

問題 ABC308F

シグネチャを決める。

abc308f :: Int    -- N
        -> Int    -- M
        -> [Int]  -- Pi
        -> [Int]  -- Li
        -> [Int]  -- Di
        -> Int    -- 答え

貪欲法でやれば済むような問題がFで出るかな？などとダメな思考にとらわれて解けなかったので解説を見た。

公式解説の解法

公式解説 by yuto1115
クーポンの使用は金額で制限があるので、安い方の商品から順に確認していく。
（高い方からやると、もっと安いものでも使えるクーポンを高い方に使ってしまい、使えないクーポンだけ残してしまって損する可能性がある。）
そして、使えるクーポンのうち、最も割引のいいものを使う。
（それより後は同額かそれ以上の商品しかないので、いつ使っても割引でもうかる額は同じだから、一番得なやつから使ってしまって構わない。）

「商品の値段で使えるクーポンは限定されるが、その中でさらに最もお得なものを選ぶ」という操作の効率的な実現に悩む。
ここで実は、ある商品について考察するときに使用許可されたクーポンは、それ以降の商品は今回のもの以上の価格のものなので、許可が取り消されるということはなく、$L_i$は忘れて$D_i$だけ覚えていればいい。（これに気付かなかった。）
そしてクーポンは、別のクーポンで$D_i = D_j$ となることもありうるので、使用可能になったクーポンの $D_i$ を多重集合で保持する、とやるのか、面倒だなぁと思ったら、どうせ探すのは最大値のみなので、割引額の降順での優先度付きキューでいけるという。（これは思いつかなかった。）

ここまで言われれば後は書くだけ。

import Data.List
import qualified Data.Heap as H

abc308f :: Int -> Int -> [Int] -> [Int] -> [Int] -> Int
abc308f n m ps ls ds =
  sum $ snd $                                  -- 3. それぞれの購入価格の総和をとる
  mapAccumL step (H.empty, sort $ zip ls ds) $ -- 2. 一番都合のいいクーポンを選んで使い、
  sort ps                                      -- 1. 価格の安い順に

type State = ( H.Heap Int       -- 割引額の降順キュー
             , [(Int,Int)])     -- まだ使えないクーポンのLi,DiをLiの昇順で

step :: State -> Int -> (State, Int)
step (h, lds) p =
  case H.uncons h1 of
    Nothing -> ((h1, lds2), p)          -- 使えるクーポンがなければ定価で買う
    Just (d, h2) -> ((h2, lds2), p + d) -- クーポンがあれば消費して値引きで買う
  where
    (lds1, lds2) = span ((p >=).fst) lds    -- 1. 今回の価格pで使用可能になったクーポンを
    h1 = H.union h $ H.fromList $ map (negate . snd) lds1 -- 2. 割引額だけをキューに追加

Data.IntMapをキュー代わりにする方が実は Data.Heap よりわずかに速かったが、大した差ではないので明快さを優先するべき。

ユーザ解説の方法

ユーザ解説 by kyopro_friends
結局貪欲法なのだけど、向きが違う。
クーポンの割引額に注目して、最もお得なものから考えていく。
そのクーポンが使える未購入の商品で最も安い物を買うのにそれを使う。
（より高い商品は、そうでないと使えないクーポンが出てきたときのために残しておく。）
最後までクーポンを使えずに残った商品は普通に買う。
という方針。

「クーポンの制限額以上の価格の、未購入の商品」を管理するには、価格をキーにした多重集合が今度こそ必要。定価をキー、商品の個数を値とするIntMapで表す。

import Data.List
import qualified Data.IntMap as IM

abc308f :: Int -> Int -> [Int] -> [Int] -> [Int] -> Int
abc308f n m ps ls ds =
  sum qs +  -- 割引で買った総額
  sum (map (uncurry (*)) $ IM.assocs pm) -- クーポンを使いそびれた商品の総額
  where
    pm0 = IM.fromListWith (+) [(p,1) | p <- ps] -- 全ての商品
    (pm,qs) = mapAccumL step pm0 $          -- 2. クーポンを使うべき商品を選んで買う
              sort $ zip (map negate ds) ls -- 1. 割引額の大きい順に

type State = IM.IntMap Int -- 未購入の商品の価格と個数

step :: State         -- 未購入の商品
     -> (Int,Int)     -- 今回のクーポン
     -> (State, Int)  -- 購入しなかったらsndは0
step pm (d, l) =
  case IM.lookupGE l pm of                          -- 価格l以上の
    Just (p, 1) -> (IM.delete p pm, p + d)          -- pの商品は最後の一つ
    Just (p, k) -> (IM.insert p (pred k) pm, p + d) -- 価格pの商品はまだ他にもある
    Nothing     -> (pm, 0)                          -- クーポンは使えなかった

優先度付きキューを使う公式解法より、IntMapを使うこちらの方が速かった。

Esprit de l'escalier

商品ごとに精算する代わりに、クーポンを使ったときにその割引額だけ合計を出しておき、$\sum P_i$ からそれを引くだけで答えは得られた。

G - Minimum Xor Pair Query

問題 ABC308G

最上位ビットについては、隣接する値とだけ比べればいいということはわかったが、それ以降のビットについては、もしかしたら隣接する値とでない方が小さくなりはしないか？と考え出してわからなくなったので解説を見た。

2進数のtrieを使うのかな？とも考えていたが、具体的なところまで詰められなかった。それはユーザ解説の方法だった。

考える

クエリ1,2では状態だけが変化し、クエリ3では結果出力だけをする。
それぞれの仕事を分担して定義し、クエリごとに駆動する形にする。

main = do
  q <- readLn
  foldM_ (\st _ -> do
    qi <- map read . words <$> getLine
    case qi of
      (1:x:_) -> return $ mode1 st x
      (2:x:_) -> return $ mode2 st x
      (3:_)   -> print (mode3 st) >> return st
    ) initial [1..q]

（実際は入力をByteStringで処理しないと時間がかかるので注意。）

公式解説の方法

黒板に書かれている数と、答えの候補であるXOR値の両方をmultisetで管理する。

import qualified Data.IntMap as IM

type State =
  ( IM.IntMap Int -- 黒板の数のmultiset
  , IM.IntMap Int -- XORの値のmultiset
  )

initial = (IM.empty, IM.empty)

insertMS x ms = IM.insertWith (+) x 1 ms
deleteMS x ms = IM.update dec x ms
  where
    dec 1 = Nothing
    dec k = Just (pred k)

値を書き込むとき、

• その値が既出なら、同じ値との間にxorの結果として0がもう一つ追加される
• その値が初出なら、両側の値の間のxorが一つ減り、両側の値と新規の値とのxorが二つ追加される（もしあれば）

mode1 (xm, am) x =
  case IM.lookup x xm of
    Just _  -> output1 0
    Nothing ->
      case (IM.lookupLT x xm, IM.lookupGT x xm) of
        (Nothing    , Nothing    ) -> output0
        (Just (x0,_), Nothing    ) -> output1 (xor x0 x)
        (Nothing    , Just (x2,_)) -> output1 (xor x2 x)
        (Just (x0,_), Just (x2,_)) -> output3 (xor x0 x2) (xor x0 x) (xor x2 x)
  where
    output am = (insertMS x xm, am)
    output0   = output am
    output1 a = output $ insertMS a am
    output3 a1 a2 a3 = output $ insertMS a3 $ insertMS a2 $ deleteMS a1 am

値を消すときは逆に、

• その値がまだ他にもあるなら、それとの間のxorの結果0を一つ減らす
• 最後の一つなら、両側との間のxorをそれぞれ減らし、両側どうしのxorを一つ追加する（もしあれば）

mode2 (xm, am) x =
  case IM.lookup x xm of
    Just 1 ->
      case (IM.lookupLT x xm, IM.lookupGT x xm) of
        (Nothing    , Nothing    ) -> output0
        (Just (x0,_), Nothing    ) -> output1 (xor x0 x)
        (Nothing    , Just (x2,_)) -> output1 (xor x2 x)
        (Just (x0,_), Just (x2,_)) -> output3 (xor x0 x2) (xor x0 x) (xor x2 x)
    Just _ -> output1 0
  where
    output am = (deleteMS x xm, am)
    output0   = output am
    output1 a = output $ deleteMS a am
    output3 a1 a2 a3 = output $ insertMS a1 $ deleteMS a2 $ deleteMS a3 am

最小値を聞かれたときは、答え候補の最小値を返す。

mode3 (_, am) = fst $ IM.findMin am

結果は 2351ms, 122MB

余計なこと

mode1とmode2において、左右の値 x0, x2 が存在するかどうかで4通りの場合分けをするのはいかにも面倒くさい。これらを Maybe で包んだまま、存在するときだけ対処するようにして格好をつけてみる。

-- insertMSとdeleteMSをMaybe対応にする
mf f Nothing  y = y
mf f (Just x) y = f x y

mode1 (xm, am) x =
  case IM.lookup x xm of
    Just _  -> output $ insertMS 0 am
    Nothing -> output3
  where
    output am = (insertMS x xm, am)
    x0 = fst <$> IM.lookupLT x xm
    x2 = fst <$> IM.lookupGT x xm
    a1 = xor x <$> x0
    a2 = xor x <$> x2
    a3 = xor <$> x0 <*> x2
    output3 = output $ mf insertMS a1 $ mf insertMS a2 $ mf deleteMS a3 am

mode2 も同様に変更した結果 2908ms, 439MBと、特にメモリが悪化した。Maybeで包まれたままでスペースリークしているのだろうと、forceやseqで候補multisetの評価を強制すると、むしろ時間が悪化してTLEするようになってしまった。
Haskell難しいな。

ユーザ解説の方法

Trieを作る。

data Trie = Node Int  -- サイズ
                 Int  -- Xの値のいずれか一つ
                 Trie -- 次の桁が0な部分木
                 Trie -- 1な部分木
          | Leaf

値の挿入、削除の処理を定義し、それぞれの動作に、答えの候補の増減に関する処理を挟んで…
解説にも

注：公式解説の方がスマートです。
本質的には公式解説と同じです。

とあるので、やめときます。