ABC 406

ABC 406 に参加しました。

A 問題

A 時 B 分 < C 時 D 分なら Yes, それ以外なら No を出力せよ。分に直しちゃったんですが、確かにタプルで比較すれば良いですね:

A はタプルで比較できる言語なら

printYn $ (a, b) >= (c, d)

でいいので。みなさん Haskell をですね
— naoya (@naoya_ito) May 17, 2025

B 問題

指定桁数以上になったら 1 に戻す掛け算に関して $\Pi_i A_i$ を求めよ。式変形すると $a b \ge U \Leftrightarrow a \ge \frac U b \Leftrightarrow a \ge \lceil \frac U a \rceil$ なので、オーバーフローを避けつつ計算できるようです:

solve :: StateT BS.ByteString IO ()
solve = do
  (!n, !k) <- ints2'
  !xs <- intsU'

  let !upper = 10 ^ k :: Int
  -- a * b >= U, a >= U/b
  let mul a b
        | a >= (upper + (b  -1)) `div` b = 1
        | otherwise = a * b
  let !res = U.foldl1' mul xs
  printBSB res

本番は多倍長整数でチートしました。 Rust の wrapping_mul は unstable のようです。

C 問題

順列 $A_i$ の連続部分列の内、単調増加、単調減少、単調増加の 3 つの部分から成るものの数を数えよ。部分列だと思ってセグ木を 4 本用意してしまいました。

隣接する二項を比較し、連長圧縮して見ると良いようですね。結構難しい:

solve :: StateT BS.ByteString IO ()
solve = do
  !n <- int'
  !xs <- U.map pred <$> intsU'
  let !xs' = U.zipWith (<) xs (U.tail xs)
  let !rle = U.fromList . map (\vec -> (U.head vec, U.length vec)) $ U.group xs'
  let !res
        | U.length rle < 3 = 0
        | otherwise = U.sum $ U.zipWith3 f rle (U.tail rle) (U.tail (U.tail rle))
        where
          f (True, !l1) (False, !_) (True, !l3) = l1 * l3
          f _ _ _ = 0
  printBSB res

D 問題

2 次元グリッドにゴミが配置されているとき、行 or 列のゴミ削除クエリに応答せよ。それぞれのゴミは高々 1 回しか削除されないことを考えると、 $O(n \log n)$ で解けます。

問題文は y 座標を x と書いていてややこしいですが、惑わされず整理します:

solve :: StateT BS.ByteString IO ()
solve = do
  (!h, !w, !n) <- ints3'
  yxs <- U.replicateM n ints2'
  q <- int'
  qs <- U.replicateM q ints2'

  -- 行ごと、列ごとにゴミを持つ
  let !rows0 = M.fromListWith S.union . map (\(!y, !x) -> (y, S.singleton x)) $ U.toList yxs
  let !cols0 = M.fromListWith S.union . map (\(!y, !x) -> (x, S.singleton y)) $ U.toList yxs

  let deleteRow y rows cols = (dn, rows', cols')
        where
          -- 削除したゴミの数
          !dn = maybe 0 S.size $ M.lookup y rows
          -- 行ごとに持ったゴミを削除する:
          !rows' = M.delete y rows
          -- 列ごとに持ったゴミも削除する:
          !cols' = case M.lookup y rows of
            Nothing -> cols
            Just xs -> foldl' (flip (M.adjust (S.delete y))) cols $ S.toList xs

  let resF = U.foldM'_ f s0 qs
        where
          s0 = (rows0, cols0)
          f (!rows, !cols) (1, !y) = do
            let (!dn, !rows', !cols') = deleteRow y rows cols
            printBSB dn
            pure (rows', cols')
          f (!rows, !cols) (2, !x) = do
            -- 行削除の関数を列削除に流用する:
            let (!dn, !cols', !rows') = deleteRow x cols rows
            printBSB dn
            pure (rows', cols')

  resF

E 問題

$\mathrm{sum} \ [x \ | \ \mathrm{popCount}(x) = k, x \le N] \bmod 998244353$ を求めよ。僕の実力が落ちているのもありますが、かなり考察が難しい問題だと思いました。基本方針は:

桁 DP で解く
配列の各スロットで個数と総和のペアを持てば、実は総和が普通に求まる

また $\mathrm{popCount}(N) = k$ の場合は答えに $N$ を加算します:

solve :: StateT BS.ByteString IO ()
solve = do
  (!n, !k) <- ints2'

  let !msb = msbOf n
  let !nBits = msb + 1

  let !res = done . U.foldl' step s0 $ U.generate nBits $ \i -> msb - i
        where
          !dn = modInt $ bool 0 n (popCount n == k)
          done :: ((Int, MyModInt), U.Vector (Int, MyModInt)) -> MyModInt
          done (!_, !res) = (dn +) . snd . U.last $ res
          s0 :: ((Int, MyModInt), U.Vector (Int, MyModInt))
          s0 = ((0 :: Int, modInt 0), U.replicate (k + 1) (0 :: Int, modInt 0))
          step :: ((Int, MyModInt), U.Vector (Int, MyModInt)) -> Int -> ((Int, MyModInt), U.Vector (Int, MyModInt))
          step ((!highPopCount, !highSum), !sofar) iBit = ((highPopCount', highSum'), sofar')
            where
              !highPopCount' = highPopCount + bool 0 1 (testBit n iBit)
              !highSum' = highSum + bool 0 (modInt (bit iBit)) (testBit n iBit)
              !sofar' = U.imap f sofar
              f iPopCount (!nAcc, !sumAcc) =
                let (!nFrom, !sumFrom)
                      | iPopCount == 0 = (0, modInt 0)
                      | otherwise = sofar G.! (iPopCount - 1)
                    -- choose `0` bit and come down from the largest number
                    !cntHigh
                      | highPopCount <= k && testBit n iBit && iPopCount == highPopCount = 1
                      | otherwise = 0
                 in (nFrom + nAcc + cntHigh, sumAcc + sumFrom + modInt nFrom * modInt (bit iBit) + modInt cntHigh * highSum)

  printBSB res

-- verification-helper: PROBLEM https://atcoder.jp/contests/abc406/tasks/abc406_e
main :: IO ()
main = runIO $ do
  t <- int'
  replicateM_ t solve

うーん汚い……。桁 DP を DFA で解くやつを導入したら綺麗になるかもしれません。

F 問題

Dynamic Tree Vertex Add Subtree Sum の簡単なやつ……と思いきや link/cut tree で TLE しました。ショックなんですが、 C++ の link/cut tree だと普通に通るようです。 ac-library-hs も C++ ほどは速くないので、 TLE しそうで辛いです。

kemuniku 氏の提出が HLD を使っていたので真似しました。片側の部分木の和が求まれば、全体の和からそれを引くことで、もう一方の部分木の和が求まります。確かにーー

solve :: StateT BS.ByteString IO ()
solve = do
  !n <- int'
  !uvs <- U.replicateM (n - 1) ints11'
  q <- int'
  qs <- U.replicateM q $ do
    int' >>= \case
      1 -> (1 :: Int,,) <$> int1' <*> int'
      2 -> (2 :: Int,,-1) <$> int1'
      _ -> error "unreachable"

  let !tree = buildSG_ n $ swapDupeU uvs
  let !hld@HLD {..} = hldOf tree
  tm <- buildVertTM hld True $ U.replicate n (Sum (1 :: Int))

  res <- (`U.mapMaybeM` qs) $ \case
    (1, !v, !dw) -> do
      modifyTM tm (+ Sum dw) v
      pure Nothing
    (2, !iEdge, !_) -> do
      let (!u, !v) = uvs G.! iEdge
      let !child = if parentHLD G.! u == v then u else v
      xWhole <- foldAllSTree (streeFTM tm)
      x1 <- foldSubtreeVertsTM tm child
      let !x2 = xWhole - x1
      pure . Just . getSum . abs $ x1 - x2

  printBSB $ unlinesBSB res

HLD と組み合わせてモノイド積を取る tree monoid については、 maspy さんの tree_monoid.hpp を参照のこと。つまりこれなんですが、 pull ではなく push ベースなのと、非可換なモノイドに対応するために 2 本のセグ木を用意します:

HLDやWavelet Matrixは内部でSegtreeは持たず、rangeの列を返すようにしたら使い勝手が良くなった
外で累積和・Fenwick Tree・Segtree・Lazy Segtreeなどから好きなものを選んで使う
— 37kt (@37kt_kyopro) May 18, 2025

Misc

`atcli`

Haskell 製の CLI で、 oj/acc のように使えるツールが登場しました。凄い。 Haskellで競プロ鉄則本を解いていくも拝見しています。