ac-library-hs 1.0.0.1 → 1.1.0.0

raw patch · 78 files changed

+9333/−709 lines, 78 filesdep +containersdep +unordered-containersPVP ok

version bump matches the API change (PVP)

Dependencies added: containers, unordered-containers

API changes (from Hackage documentation)

Files

• CHANGELOG.md +11/−3
• README.md +2/−2
• ac-library-hs.cabal +45/−3
• benchmarks/Bench/AddMod.hs +1/−1
• benchmarks/Bench/Matrix.hs +87/−0
• benchmarks/Bench/ModInt.hs +23/−0
• benchmarks/BenchLib/Matrix.hs +394/−0
• benchmarks/BenchLib/MulMod/BarrettWideWord.hs +2/−1
• benchmarks/BenchLib/MulMod/Montgomery.hs +3/−3
• benchmarks/Main.hs +8/−5
• src/AtCoder/Convolution.hs +2/−2
• src/AtCoder/Dsu.hs +16/−16
• src/AtCoder/Extra/Bisect.hs +297/−0
• src/AtCoder/Extra/Graph.hs +147/−0
• src/AtCoder/Extra/HashMap.hs +346/−0
• src/AtCoder/Extra/IntMap.hs +338/−0
• src/AtCoder/Extra/IntSet.hs +391/−0
• src/AtCoder/Extra/IntervalMap.hs +462/−0
• src/AtCoder/Extra/Math.hs +34/−18
• src/AtCoder/Extra/Monoid.hs +23/−12
• src/AtCoder/Extra/Monoid/Affine1.hs +40/−14
• src/AtCoder/Extra/Monoid/Mat2x2.hs +202/−0
• src/AtCoder/Extra/Monoid/RangeAdd.hs +38/−22
• src/AtCoder/Extra/Monoid/RangeAddId.hs +0/−95
• src/AtCoder/Extra/Monoid/RangeSet.hs +18/−10
• src/AtCoder/Extra/Monoid/RangeSetId.hs +0/−124
• src/AtCoder/Extra/Monoid/RollingHash.hs +121/−0
• src/AtCoder/Extra/Monoid/V2.hs +81/−0
• src/AtCoder/Extra/MultiSet.hs +280/−0
• src/AtCoder/Extra/Pdsu.hs +306/−0
• src/AtCoder/Extra/Semigroup/Matrix.hs +264/−0
• src/AtCoder/Extra/Semigroup/Permutation.hs +114/−0
• src/AtCoder/Extra/Tree.hs +207/−0
• src/AtCoder/Extra/Tree/Hld.hs +512/−0
• src/AtCoder/Extra/Tree/TreeMonoid.hs +297/−0
• src/AtCoder/Extra/WaveletMatrix.hs +407/−0
• src/AtCoder/Extra/WaveletMatrix/BitVector.hs +179/−0
• src/AtCoder/Extra/WaveletMatrix/Raw.hs +651/−0
• src/AtCoder/Extra/WaveletMatrix2d.hs +277/−0
• src/AtCoder/FenwickTree.hs +10/−10
• src/AtCoder/Internal/Assert.hs +1/−1
• src/AtCoder/Internal/Barrett.hs +6/−6
• src/AtCoder/Internal/Buffer.hs +54/−59
• src/AtCoder/Internal/Convolution.hs +1/−1
• src/AtCoder/Internal/Csr.hs +47/−16
• src/AtCoder/Internal/GrowVec.hs +40/−40
• src/AtCoder/Internal/Math.hs +19/−3
• src/AtCoder/Internal/MinHeap.hs +18/−18
• src/AtCoder/Internal/Queue.hs +37/−35
• src/AtCoder/Internal/Scc.hs +55/−21
• src/AtCoder/Internal/String.hs +33/−4
• src/AtCoder/LazySegTree.hs +40/−17
• src/AtCoder/Math.hs +23/−4
• src/AtCoder/MaxFlow.hs +82/−12
• src/AtCoder/MinCostFlow.hs +36/−4
• src/AtCoder/ModInt.hs +40/−37
• src/AtCoder/Scc.hs +7/−1
• src/AtCoder/SegTree.hs +70/−13
• src/AtCoder/String.hs +37/−20
• src/AtCoder/TwoSat.hs +1/−1
• test/Main.hs +29/−2
• test/Tests/Extra/Bisect.hs +114/−0
• test/Tests/Extra/Graph.hs +52/−0
• test/Tests/Extra/HashMap.hs +171/−0
• test/Tests/Extra/IntMap.hs +176/−0
• test/Tests/Extra/IntSet.hs +150/−0
• test/Tests/Extra/IntervalMap.hs +209/−0
• test/Tests/Extra/Monoid.hs +97/−49
• test/Tests/Extra/MultiSet.hs +218/−0
• test/Tests/Extra/Semigroup/Matrix.hs +43/−0
• test/Tests/Extra/Semigroup/Permutation.hs +71/−0
• test/Tests/Extra/WaveletMatrix.hs +237/−0
• test/Tests/Extra/WaveletMatrix/BitVector.hs +84/−0
• test/Tests/Extra/WaveletMatrix/Raw.hs +271/−0
• test/Tests/Extra/WaveletMatrix2d.hs +119/−0
• test/Tests/MaxFlow.hs +1/−1
• test/Tests/Scc.hs +0/−2
• test/Tests/Util.hs +8/−1

@@ -1,8 +1,16 @@ # Revision history for acl-hs +## 1.1.0.0 -- Jan 2025++- Removed `RangeSetId` and `RangeAddId` from `AtCoder.Extra.Monoid`.+- Implemented `SegAct` for `RangeSet`, `RangeAdd` and `Max`, `Min`.+- Added `segActWithLength` to `SegAct`.+- Added `build1` to `AtCoder.Internal.Csr`.+- Added a bunch of extra modules.+ ## 1.0.0.0 -- Dec 2024 -* First version.-* ACL-compatible modules.-* Extra module of `Math` (binary exponentiation) and `Monoid` (`SegAct` instances).+- First version.+- Added ACL-compatible modules.+- Added Extra module of `Math` (binary exponentiation) and `Monoid` (`SegAct` instances). 

@@ -5,8 +5,8 @@ ## Notes  - The library is mainly for AtCoder and only GHC 9.8.4 is guaranteed to be supported.-- The `Internal` module is unstable and does not follow PVP.-- The `Extra` module provides with extra contents not in the original `ac-library`. Note that they're less tested.+- Functions primarily use half-open interval \([l, r)\).+- The extra module contains additional utilities beyond the original C++ library.  ## Usage 

@@ -4,12 +4,15 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.0.0.1+version:         1.1.0.0 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive   programming on [AtCoder](https://atcoder.jp/). +  - Functions primarily use half-open interval \([l, r)\).+  - The extra module contains additional utilities beyond the original C++ library.+ category:        Algorithms, Data Structures license:         CC0-1.0 license-file:    LICENSE@@ -70,13 +73,31 @@   exposed-modules:     AtCoder.Convolution     AtCoder.Dsu+    AtCoder.Extra.Bisect+    AtCoder.Extra.HashMap+    AtCoder.Extra.IntervalMap+    AtCoder.Extra.IntMap+    AtCoder.Extra.IntSet+    AtCoder.Extra.Graph+    AtCoder.Extra.Tree+    AtCoder.Extra.Tree.Hld+    AtCoder.Extra.Tree.TreeMonoid     AtCoder.Extra.Math     AtCoder.Extra.Monoid     AtCoder.Extra.Monoid.Affine1+    AtCoder.Extra.Monoid.Mat2x2     AtCoder.Extra.Monoid.RangeAdd-    AtCoder.Extra.Monoid.RangeAddId     AtCoder.Extra.Monoid.RangeSet-    AtCoder.Extra.Monoid.RangeSetId+    AtCoder.Extra.Monoid.RollingHash+    AtCoder.Extra.Monoid.V2+    AtCoder.Extra.MultiSet+    AtCoder.Extra.Semigroup.Matrix+    AtCoder.Extra.Semigroup.Permutation+    AtCoder.Extra.Pdsu+    AtCoder.Extra.WaveletMatrix+    AtCoder.Extra.WaveletMatrix.BitVector+    AtCoder.Extra.WaveletMatrix.Raw+    AtCoder.Extra.WaveletMatrix2d     AtCoder.FenwickTree     AtCoder.Internal.Assert     AtCoder.Internal.Barrett@@ -118,8 +139,21 @@   other-modules:     Tests.Convolution     Tests.Dsu+    Tests.Extra.Bisect+    Tests.Extra.Graph+    Tests.Extra.HashMap+    Tests.Extra.IntervalMap+    Tests.Extra.IntMap+    Tests.Extra.IntSet     Tests.Extra.Math     Tests.Extra.Monoid+    Tests.Extra.MultiSet+    Tests.Extra.Semigroup.Matrix+    Tests.Extra.Semigroup.Permutation+    Tests.Extra.WaveletMatrix+    Tests.Extra.WaveletMatrix.BitVector+    Tests.Extra.WaveletMatrix.Raw+    Tests.Extra.WaveletMatrix2d     Tests.FenwickTree     Tests.Internal.Bit     Tests.Internal.Buffer@@ -146,6 +180,7 @@   main-is:        Main.hs   build-depends:     , ac-library-hs+    , containers     , hspec     , mtl     , QuickCheck@@ -157,6 +192,7 @@     , tasty-quickcheck     , tasty-rerun     , transformers+    , unordered-containers  benchmark ac-library-hs-benchmark   import:         warnings@@ -167,8 +203,11 @@   hs-source-dirs: benchmarks   other-modules:     Bench.AddMod+    Bench.Matrix+    Bench.ModInt     Bench.PowMod     BenchLib.AddMod+    BenchLib.Matrix     BenchLib.ModInt.ModIntNats     BenchLib.ModInt.Modulus     BenchLib.MulMod.Barrett64@@ -180,8 +219,10 @@     , ac-library-hs     , base     , criterion+    , mtl     , random     , tagged+    , transformers  test-suite benchlib-test   import:         warnings@@ -204,6 +245,7 @@     , hspec     , mtl     , QuickCheck+    -- , quickcheck-instances     , random     , tagged     , tasty

@@ -6,7 +6,7 @@ import BenchLib.AddMod import Criterion import Data.Vector.Unboxed qualified as VU-import Data.Word (Word32, Word64)+import Data.Word (Word32) import System.Random  n :: Int

@@ -0,0 +1,87 @@+module Bench.Matrix (benches) where++import AtCoder.ModInt qualified as M+import AtCoder.Extra.Semigroup.Matrix qualified as ACMAT+import AtCoder.Extra.Math qualified as ACEM+import BenchLib.Matrix qualified as Mat+import Control.Monad.State.Class (MonadState, state)+import Control.Monad.Trans.State.Strict (evalState, runState)+import Criterion+import Data.Semigroup (Semigroup (..))+import Data.Vector qualified as V+import Data.Vector.Unboxed qualified as VU+import System.Random++testSize :: Int+testSize = 10000++m :: Int+m = 998244353++randomMatrix :: (MonadState StdGen m) => m (Mat.Matrix Int, Int)+randomMatrix = do+  h <- state $ uniformR (1, 16)+  nextMatrix h++nextMatrix :: (MonadState StdGen m) => Int -> m (Mat.Matrix Int, Int)+nextMatrix h = do+  w <- state $ uniformR (1, 16)+  vec <- VU.replicateM (h * w) $ state (uniformR (0, m - 1))+  pure (Mat.new h w vec, w)++randomSquareACLMatrix :: (MonadState StdGen m) => Int -> m (ACMAT.Matrix Int)+randomSquareACLMatrix n = do+  vec <- VU.replicateM (n * n) $ state (uniformR (0, m - 1))+  pure $ ACMAT.new n n vec++benches :: Benchmark+benches =+  bgroup+    "Matrix"+    [ bench "mul1" $ whnf (V.foldl1' Mat.mul1) randomMatrixInput,+      bench "mul2" $ whnf (V.foldl1' Mat.mul2) randomMatrixInput,+      bench "mul3_1" $ whnf (V.foldl1' Mat.mul3_1) randomMatrixInput,+      bench "mul3_2" $ whnf (V.foldl1' Mat.mul3_2) randomMatrixInput,+      bench "mul3_3" $ whnf (V.foldl1' Mat.mul3_3) randomMatrixInput,+      bench "mulMod1" $ whnf (V.foldl1' (Mat.mulMod1 m)) randomMatrixInput,+      bench "mulMod2" $ whnf (V.foldl1' (Mat.mulMod2 m)) randomMatrixInput,+      bench "mulMod3" $ whnf (V.foldl1' (Mat.mulMod3 m)) randomMatrixInput,+      bench "mulMod4" $ whnf (V.foldl1' (Mat.mulMod4 m)) randomMatrixInput,+      bench "mulMod5" $ whnf (V.foldl1' (Mat.mulMod5 m)) randomMatrixInput,+      bench "mulMint1" $ whnf (V.foldl1' Mat.mulMint1) randomMintMatrixInput,+      bench "mulMint2" $ whnf (V.foldl1' Mat.mulMint2) randomMintMatrixInput,+      bench "mulMint3" $ whnf (V.foldl1' Mat.mulMint3) randomMintMatrixInput,+      -- mul (ACL)+      bench "mul_ACL" $ whnf (V.foldl1' ACMAT.mul) randomMatrixInputACL,+      bench "mulMod_ACL" $ whnf (V.foldl1' (ACMAT.mulMod m)) randomMatrixInputACL,+      -- pow mod (ACL only)+      bench "powMod_ACL" $ whnf (VU.foldl' (flip (ACMAT.powMod m)) squareMat) randomVec,+      bench "powMintACL_stimes" $ whnf (VU.foldl' (flip stimes) squareMatMint) randomVec,+      bench "powMintACL_stimes'" $ whnf (VU.foldl' (flip ACEM.stimes') squareMatMint) randomVec,+      bench "powMintACL_powMint" $ whnf (VU.foldl' (flip ACMAT.powMint) squareMatMint) randomVec+    ]+  where+    -- Bench matrix+    randomMatrixInput :: V.Vector (Mat.Matrix Int)+    !randomMatrixInput = evalState (V.unfoldrExactNM testSize nextMatrix (Mat.wM mat0)) gen0+      where+        ((!mat0, !_), !gen0) = runState randomMatrix $ mkStdGen 123456789++    randomMintMatrixInput :: V.Vector (Mat.Matrix (M.ModInt 998244353))+    !randomMintMatrixInput = V.map (Mat.map M.new) randomMatrixInput++    -- ACL matrix+    randomMatrixInputACL :: V.Vector (ACMAT.Matrix Int)+    !randomMatrixInputACL = V.map (\mat -> ACMAT.new (Mat.hM mat) (Mat.wM mat) (Mat.vecM mat)) randomMatrixInput++    squareMat :: ACMAT.Matrix Int+    !squareMat = evalState (randomSquareACLMatrix 17) $ mkStdGen 123456789++    squareMatMint ::ACMAT.Matrix (M.ModInt 998244353)+    !squareMatMint = ACMAT.map M.new squareMat++    -- non-zero random vector+    randomVec :: VU.Vector Int+    !randomVec =+      VU.map ((+ 1) . fromIntegral) $+        VU.unfoldrExactN 100 (genWord64R (m - 2)) (mkStdGen 123456789)

@@ -0,0 +1,23 @@+module Bench.ModInt (benches) where++import AtCoder.ModInt qualified as M+import BenchLib.PowMod qualified as PowMod+import BenchLib.ModInt.ModIntNats qualified as MN+import Criterion+import Data.Vector.Unboxed qualified as VU+import System.Random++benches :: Benchmark+benches =+  bgroup+    "modInt"+    [++    ]+  where+    n = 10000+    randomVec :: VU.Vector Int+    randomVec =+      VU.map fromIntegral $+        VU.unfoldrExactN n (genWord64R (998244383 - 2)) (mkStdGen 123456789)+

@@ -0,0 +1,394 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE RecordWildCards #-}++module BenchLib.Matrix+  ( Matrix (..),+    new,+    map,+    mulToCol,+    mulToColModInt,+    mul1,+    mul2,+    mul3_1,+    mul3_2,+    mul3_3,+    mulMod1,+    mulMod2,+    mulMod3,+    mulMod4,+    mulMod5,+    mulMint1,+    mulMint2,+    mulMint3,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Internal.Barrett qualified as BT+import AtCoder.ModInt qualified as M+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU+import Data.Word (Word64)+import GHC.Exts (Proxy#, proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal, natVal')+import Prelude hiding (map)++data Matrix a = Matrix+  { hM :: {-# UNPACK #-} !Int,+    wM :: {-# UNPACK #-} !Int,+    vecM :: !(VU.Vector a)+  }+  deriving (Show, Eq)++type Col a = VU.Vector a++{-# INLINE new #-}+new :: (HasCallStack, VU.Unbox a) => Int -> Int -> VU.Vector a -> Matrix a+new h w vec+  | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch"+  | otherwise = Matrix h w vec++{-# INLINE map #-}+map :: (VU.Unbox a, VU.Unbox b) => (a -> b) -> Matrix a -> Matrix b+map f Matrix {..} = Matrix hM wM $ VU.map f vecM++{-# INLINE mulToCol #-}+mulToCol :: (Num a, VU.Unbox a) => Matrix a -> Col a -> Col a+mulToCol Matrix {..} !col = VU.convert $ V.map (VU.sum . VU.zipWith (*) col) rows+  where+    !n = VU.length col+    !_ = ACIA.runtimeAssert (n == wM) "AtCoder.Extra.Matrix.mulToCol: size mismatch"+    rows = V.unfoldrExactN hM (VU.splitAt wM) vecM++{-# INLINE mulToColModInt #-}+mulToColModInt :: forall m. (KnownNat m) => Matrix (M.ModInt m) -> Col (M.ModInt m) -> Col (M.ModInt m)+mulToColModInt Matrix {..} !col = VU.convert $ V.map (VU.foldl' (+) (M.unsafeNew 0) . VU.zipWith mulMod col) rows+  where+    !_ = ACIA.runtimeAssert (VU.length col == wM) "AtCoder.Extra.Matrix.mulToColModInt: size mismatch"+    !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m))+    rows = V.unfoldrExactN hM (VU.splitAt wM) vecM+    mulMod (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)++{-# INLINE mul1 #-}+mul1 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul1 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.sum $ VU.zipWith (*) (rows1 VG.! row) (cols2 VG.! col)+    h = hM a+    w = wM a+    vecA = vecM a+    h' = hM b+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"+    rows1 = V.unfoldrExactN h (VU.splitAt w) vecA+    cols2 = V.generate w' $ \col -> VU.unfoldrExactN h' (\i -> (VG.unsafeIndex vecB i, i + w')) col++{-# INLINE mul2 #-}+mul2 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul2 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (rows1 VG.! row)+    h = hM a+    w = wM a+    vecA = vecM a+    h' = hM b+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"+    rows1 = V.unfoldrExactN h (VU.splitAt w) vecA++{-# INLINE mul3_1 #-}+mul3_1 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul3_1 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mul3_2 #-}+mul3_2 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul3_2 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !y = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)+           in if col + 1 >= w'+                then (y, (row + 1, 0))+                else (y, (row, col + 1))+      )+      (0, 0)+  where+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mul3_3 #-}+mul3_3 :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul3_3 !a !b = Matrix h w' $ VU.generate (h * w') $ \i ->+  let (!row, !col) = i `quotRem` w'+   in VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)+  where+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod1 #-}+mulMod1 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod1 !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    addMod x y = (x + y) `mod` m+    mulMod x y = (x * y) `mod` m+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod2 #-}+mulMod2 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod2 !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    -- very slow+    addMod x y+      | x + y >= m = x + y - m+      | otherwise = x + y+    mulMod x y = (x * y) `mod` m+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod3 #-}+mulMod3 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod3 !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral m+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    addMod x y = (x + y) `mod` m+    mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod4 #-}+mulMod4 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod4 !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral m+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    addMod x y = (x + y) `rem` m+    mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMod5 #-}+mulMod5 :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod5 !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral m+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    addMod x y+      | x + y >= m = x + y - m+      | otherwise = x + y+    mulMod x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMint1 #-}+mulMint1 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMint1 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f :: Int -> Int -> M.ModInt a+    f row col = VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    mulMod :: M.ModInt a -> M.ModInt a -> M.ModInt a+    mulMod = (*)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++{-# INLINE mulMint2 #-}+mulMint2 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMint2 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))+    f :: Int -> Int -> M.ModInt a+    f row col = VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    mulMod :: M.ModInt a -> M.ModInt a -> M.ModInt a+    mulMod (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++-- REMARK: This is very unsafe in that it can overflow (mod^2 * n)+{-# INLINE mulMint3 #-}+mulMint3 :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMint3 !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))+    f :: Int -> Int -> M.ModInt a+    f row col = M.new64 . VU.sum $ VU.imap (\iRow x -> mulMod x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    mulMod :: M.ModInt a -> M.ModInt a -> Word64+    mulMod (M.ModInt x) (M.ModInt y) = BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"

@@ -24,10 +24,11 @@     imBarrett :: {-# UNPACK #-} !Word64   } --- | Creates barret reduction for modulus \(m\).+-- | Creates a `Barrett` for modulus value \(m\) of type `Word32`. new32 :: Word32 -> Barrett new32 m = Barrett m $ maxBound @Word64 `div` (fromIntegral m :: Word64) + 1 +-- | Creates a `Barrett` for modulus value \(m\) of type `Word64`. new64 :: Word64 -> Barrett new64 m = Barrett (fromIntegral m) $ maxBound @Word64 `div` m + 1 

@@ -17,8 +17,8 @@ import Data.Bits (bit, (!>>.)) import Data.Word (Word32, Word64) --- | Fast modular multiplication by Montgomery multiplication. The modulus value has to be odd--- to be fast.+-- | Fast modular multiplication by Montgomery multiplication. The modulus value must be odd for+-- the speed. data Montgomery = Montgomery   { mMontgomery :: {-# UNPACK #-} !Word64,     -- | R2 == (2^64) % MOD;@@ -27,7 +27,7 @@     negInvMontgomery :: {-# UNPACK #-} !Word32   } --- | Creates `Montgomery` for modulus @m@.+-- | Creates a `Montgomery` for modulus @m@. new :: Word64 -> Montgomery new m =   let !negInv = inner 0 0 1 0

@@ -3,7 +3,9 @@ module Main where  import Bench.AddMod qualified+import Bench.Matrix qualified import Bench.MulMod qualified+import Bench.ModInt qualified import Bench.PowMod qualified import Criterion.Main @@ -13,9 +15,10 @@ main :: IO () main =   defaultMain-  -- TODO: generate graph by benchmark group?-    [ -- Bench.MulMod.benches32,-      -- Bench.MulMod.benches64,-      -- Bench.AddMod.benches-      Bench.PowMod.benches+    -- TODO: generate criterion graph by benchmark group?+    [ Bench.MulMod.benches,+      Bench.ModInt.benches,+      Bench.AddMod.benches,+      Bench.PowMod.benches,+      Bench.Matrix.benches     ]

@@ -14,7 +14,7 @@ -- >>> import Data.Proxy (Proxy) -- >>> import Data.Vector.Unboxed qualified as VU ----- It's handy to define specific items for interested modulus values:+-- Define specific modulus items: -- -- >>> type Mint = M.ModInt998244353 -- >>> let modInt :: Int -> Mint; modInt = M.new@@ -26,7 +26,7 @@ -- >>> C.convolution a b -- [5,16,34,60,61,52,32] ----- You can also target any @'Integral' a@ (with some runtime overhead for conversion to modint):+-- You can also target any @'Integral' a@ with `convolutionRaw`: -- -- >>> let a = VU.fromList @Int [1, 2, 3, 4] -- >>> let b = VU.fromList @Int [5, 6, 7, 8]

@@ -1,17 +1,17 @@ {-# LANGUAGE RecordWildCards #-} --- | Disjoint set union, also known as Union-Find tree. It processes the following queries in+-- | A disjoint set union, also known as a Union-Find tree. It processes the following queries in -- amortized \(O(\alpha(n))\) time. ----- - Edge addition--- - Deciding whether given two vertices are in the same connected component+-- - Edge addition (`merge`)+-- - Deciding whether given two vertices are in the same connected component (`same`) ----- Each connected component internally has a representative vertex. When two connected components--- are merged by edge addition, one of the two representatives of these connected components--- becomes the representative of the new connected component.+-- Each connected component internally has a representative vertex (`leader`). When two connected+-- components are merged by edge addition (`merge`), one of the two representatives of these+-- connected components becomes the representative (`leader`) of the new connected component. -- -- ==== __Example__--- Create DSU:+-- Create a `Dsu` with four vertices: -- -- >>> import AtCoder.Dsu qualified as Dsu -- >>> dsu <- Dsu.new 4   -- 0 1 2 3@@ -49,15 +49,15 @@     -- * Constructor     new, -    -- * Merge+    -- * Merging     merge,     merge_,      -- * Leader     leader,-    same,      -- * Component information+    same,     size,     groups,   )@@ -73,13 +73,13 @@ import Data.Vector.Unboxed.Mutable qualified as VUM import GHC.Stack (HasCallStack) --- | Disjoint set union. Akso known as Union-Find tree.+-- | A disjoint set union. Akso known as Union-Find tree. -- -- @since 1.0.0.0 data Dsu s = Dsu-  { -- | 1.0.0 The number of nodes.+  { -- | The number of nodes.     ---    -- @since+    -- @since 1.0.0.0     nDsu :: {-# UNPACK #-} !Int,     -- | For root (leader) nodes it stores their size as a negative number. For child nodes it     -- stores their parent node index.@@ -103,8 +103,8 @@       pure Dsu {..}   | otherwise = error $ "new: given negative size (`" ++ show nDsu ++ "`)" --- | Adds an edge @(a, b)@. the vertices @a@ and @b@ were in the same connected component, it--- returns the representative of this connected component. Otherwise, it returns the+-- | Adds an edge \((a, b)\). If the vertices \(a\) and \(b\) are in the same connected component, it+-- returns the representative (`leader`) of this connected component. Otherwise, it returns the -- representative of the new connected component. -- -- ==== Constraints@@ -134,7 +134,7 @@       VGM.modify parentOrSizeDsu (+ sizeY) x       pure x --- | `merge` with return value discarded.+-- | `merge` with the return value discarded. -- -- ==== Constraints -- - \(0 \leq a < n\)@@ -207,7 +207,7 @@   sizeLa <- VGM.read parentOrSizeDsu la   pure (-sizeLa) --- | Divides the graph into connected components and returns the list of them.+-- | Divides the graph into connected components and returns the vector of them. -- -- More precisely, it returns a vector of the "vector of the vertices in a connected component". -- Both of the orders of the connected components and the vertices are undefined.

@@ -0,0 +1,297 @@+-- | Bisection methods and binary search functions. They partition a half-open interval \([l, r)\)+-- into two and return either the left or the right point of the boundary.+--+-- @+-- Y Y Y Y Y N N N N N      Y: user predicate holds+-- --------* *---------> X  N: user predicate does not hold+--         L R              L, R: left, right point of the boundary+-- @+--+-- ==== __Example__+-- Perform index compression:+--+-- >>> import AtCoder.Extra.Bisect+-- >>> import Data.Maybe (fromJust)+-- >>> import Data.Vector.Algorithms.Intro qualified as VAI+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList ([0, 20, 10, 40, 30] :: [Int])+-- >>> let dict = VU.uniq $ VU.modify VAI.sort xs+-- >>> VU.map (fromJust . lowerBound dict) xs+-- [0,2,1,4,3]+--+-- @since 1.1.0.0+module AtCoder.Extra.Bisect+  ( -- * C++-like binary search+    lowerBound,+    lowerBoundIn,+    upperBound,+    upperBoundIn,++    -- * Generic bisection method+    bisectL,+    bisectLM,+    bisectR,+    bisectRM,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Data.Functor ((<&>))+import Data.Functor.Identity+import Data.Vector.Generic qualified as VG+import GHC.Stack (HasCallStack)++-- | \(O(\log n)\) Bisection method implementation. Works on a half-open interfal \([l, r)\) .+--+-- @since 1.1.0.0+{-# INLINE bisectLImpl #-}+bisectLImpl :: (HasCallStack, Monad m) => (Int -> m Bool) -> Int -> Int -> m Int+bisectLImpl p l0 = inner (l0 - 1)+  where+    inner l r+      | l + 1 == r = pure l+      | otherwise =+          p mid >>= \case+            True -> inner mid r+            False -> inner l mid+      where+        mid = (l + r) `div` 2++-- | \(O(\log n)\) Bisection method implementation. Works on a half-open interfal \([l, r)\) .+--+-- @since 1.1.0.0+{-# INLINE bisectRImpl #-}+bisectRImpl :: (HasCallStack, Monad m) => (Int -> m Bool) -> Int -> Int -> m Int+bisectRImpl p l = ((+ 1) <$>) . bisectLImpl p l++-- | \(O(\log n)\) Returns the index of the first element \(x\) in the vector such that+-- \(x \ge x_0\), or `Nothing` if no such element exists.+--+-- @+-- Y Y Y Y Y N N N N N      Y: (< x0)+-- --------- *---------> X  N: (>= x0)+--           R              R: returning point+-- @+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [1, 1, 2, 2, 4, 4]+-- >>> lowerBound xs 1+-- Just 0+--+-- >>> lowerBound xs 2+-- Just 2+--+-- >>> lowerBound xs 3+-- Just 4+--+-- >>> lowerBound xs 4+-- Just 4+--+-- Out of range values:+--+-- >>> lowerBound xs 0+-- Just 0+--+-- >>> lowerBound xs 5+-- Nothing+--+-- @since 1.1.0.0+{-# INLINE lowerBound #-}+lowerBound :: (HasCallStack, VG.Vector v a, Ord a) => v a -> a -> Maybe Int+lowerBound vec = lowerBoundIn 0 (VG.length vec) vec++-- | \(O(\log n)\) Computes the `lowerBound` for a slice of a vector within the interval \([l, r)\).+--+-- - The user predicate evaluates indices in \([l, r)\).+-- - The interval \([l, r)\) is silently clamped to ensure it remains within the bounds \([0, n)\).+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40]+-- >>> --                            *---*---*+-- >>> lowerBoundIn 2 5 xs 10+-- Just 2+--+-- >>> lowerBoundIn 2 5 xs 20+-- Just 2+--+-- >>> lowerBoundIn 2 5 xs 30+-- Just 4+--+-- >>> lowerBoundIn 2 5 xs 40+-- Just 4+--+-- >>> lowerBoundIn 2 5 xs 50+-- Nothing+--+-- @since 1.1.0.0+{-# INLINE lowerBoundIn #-}+lowerBoundIn :: (VG.Vector v a, Ord a) => Int -> Int -> v a -> a -> Maybe Int+lowerBoundIn l_ r_ vec target+  | ACIA.testInterval l r (VG.length vec) = bisectR l r $ \i -> VG.unsafeIndex vec i < target+  | otherwise = Nothing+  where+    -- clamp+    l = max 0 l_+    r = min (VG.length vec) r_++-- | \(O(\log n)\) Returns the index of the first element \(x\) in the vector such that+-- \(x \gt x_0\), or `Nothing` if no such element exists.+--+-- @+-- Y Y Y Y Y N N N N N      Y: (<= x0)+-- --------- *---------> X  N: (> x0)+--           R              R: returning point+-- @+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40]+-- >>> upperBound xs 10+-- Just 2+--+-- >>> upperBound xs 20+-- Just 4+--+-- >>> upperBound xs 30+-- Just 4+--+-- >>> upperBound xs 40+-- Nothing+--+-- Out of range values:+--+-- >>> upperBound xs 0+-- Just 0+--+-- >>> upperBound xs 50+-- Nothing+--+-- @since 1.1.0.0+{-# INLINE upperBound #-}+upperBound :: (HasCallStack, VG.Vector v a, Ord a) => v a -> a -> Maybe Int+upperBound vec = upperBoundIn 0 (VG.length vec) vec++-- | \(O(\log n)\) Computes the `upperBound` for a slice of a vector within the interval \([l, r)\).+--+-- - The user predicate evaluates indices in \([l, r)\).+-- - The interval \([l, r)\) is silently clamped to ensure it remains within the bounds \([0, n)\).+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [10, 10, 20, 20, 40, 40]+-- >>> --                            *---*---*+-- >>> upperBoundIn 2 5 xs 0+-- Just 2+--+-- >>> upperBoundIn 2 5 xs 10+-- Just 2+--+-- >>> upperBoundIn 2 5 xs 20+-- Just 4+--+-- >>> upperBoundIn 2 5 xs 30+-- Just 4+--+-- >>> upperBoundIn 2 5 xs 40+-- Nothing+--+-- >>> upperBoundIn 2 5 xs 50+-- Nothing+--+-- @since 1.1.0.0+{-# INLINE upperBoundIn #-}+upperBoundIn :: (VG.Vector v a, Ord a) => Int -> Int -> v a -> a -> Maybe Int+upperBoundIn l_ r_ vec target+  | ACIA.testInterval l r (VG.length vec) = bisectR l r $ \i -> VG.unsafeIndex vec i <= target+  | otherwise = Nothing+  where+    -- clamp+    l = max 0 l_+    r = min (VG.length vec) r_++-- | \(O(\log n)\) Applies the bisection method on a half-open interval \([l, r)\) and returns the+-- left boundary point, or `Nothing` if no such point exists.+--+-- @+-- Y Y Y Y Y N N N N N      Y: user predicate holds+-- --------* ----------> X  N: user predicate does not hold+--         L                L: the left boundary point returned+-- @+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [10, 10, 20, 20, 30, 30]+-- >>> let n = VU.length xs+-- >>> bisectL 0 n ((<= 20) . (xs VU.!))+-- Just 3+--+-- >>> bisectL 0 n ((<= 0) . (xs VU.!))+-- Nothing+--+-- >>> bisectL 0 n ((<= 100) . (xs VU.!))+-- Just 5+--+-- >>> bisectL 0 3 ((<= 20) . (xs VU.!))+-- Just 2+--+-- @since 1.1.0.0+{-# INLINE bisectL #-}+bisectL :: (HasCallStack) => Int -> Int -> (Int -> Bool) -> Maybe Int+bisectL l r p = runIdentity $ bisectLM l r (pure . p)++-- | \(O(\log n)\) Monadic variant of `bisectL`.+--+-- @since 1.1.0.0+{-# INLINE bisectLM #-}+bisectLM :: (HasCallStack, Monad m) => Int -> Int -> (Int -> m Bool) -> m (Maybe Int)+bisectLM l r p+  | l >= r = pure Nothing+  | otherwise =+      bisectLImpl p l r <&> \case+        i | i == (l - 1) -> Nothing+        i -> Just i++-- | \(O(\log n)\) Applies the bisection method on a half-open interval \([l, r)\) and returns the+-- right boundary point, or `Nothing` if no such point exists.+--+--+-- @+-- Y Y Y Y Y N N N N N      Y: user predicate holds+-- --------- *---------> X  N: user predicate does not hold+--           R              R: the right boundary point returned+-- @+--+-- ==== __Example__+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let xs = VU.fromList [10, 10, 20, 20, 30, 30]+-- >>> let n = VU.length xs+-- >>> bisectR 0 n ((<= 20) . (xs VU.!))+-- Just 4+--+-- >>> bisectR 0 n ((<= 0) . (xs VU.!))+-- Just 0+--+-- >>> bisectR 0 n ((<= 100) . (xs VU.!))+-- Nothing+--+-- >>> bisectR 0 4 ((<= 20) . (xs VU.!))+-- Nothing+--+-- @since 1.1.0.0+{-# INLINE bisectR #-}+bisectR :: (HasCallStack) => Int -> Int -> (Int -> Bool) -> Maybe Int+bisectR l r p = runIdentity $ bisectRM l r (pure . p)++-- | \(O(\log n)\) Monadic variant of `bisectR`.+--+-- @since 1.1.0.0+{-# INLINE bisectRM #-}+bisectRM :: (HasCallStack, Monad m) => Int -> Int -> (Int -> m Bool) -> m (Maybe Int)+bisectRM l r p+  | l >= r = pure Nothing+  | otherwise =+      bisectRImpl p l r <&> \case+        i | i == r -> Nothing+        i -> Just i

@@ -0,0 +1,147 @@+-- | Re-export of the @Csr@ module and generic graph search functions.+--+-- @since 1.1.0.0+module AtCoder.Extra.Graph+  ( -- * Re-export of CSR++    -- | The `Csr.Csr` data type and all the functions such as `build` or `adj` are re-exported.+    module Csr,++    -- * CSR helpers+    swapDupe,+    swapDupe',+    scc,++    -- * Graph search+    topSort,+  )+where++import AtCoder.Extra.IntSet qualified as IS+import AtCoder.Internal.Buffer qualified as B+import AtCoder.Internal.Csr as Csr+import AtCoder.Internal.Scc qualified as ACISCC+import Control.Monad (when)+import Control.Monad.ST (runST)+import Data.Foldable (for_)+import Data.Vector qualified as V+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM++-- | \(O(n)\) Converts non-directed edges into directional edges. This is a convenient function for+-- making an input to `build`.+--+-- ==== __Example__+-- `swapDupe` duplicates each edge reversing the direction:+--+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())]+-- [(0,1,()),(1,0,()),(1,2,()),(2,1,())]+--+-- Create a non-directed graph:+--+-- >>> let gr = Gr.build 3 . Gr.swapDupe $ VU.fromList [(0, 1, ()), (1, 2, ())]+-- >>> gr `Gr.adj` 0+-- [1]+--+-- >>> gr `Gr.adj` 1+-- [0,2]+--+-- >>> gr `Gr.adj` 2+-- [1]+--+-- @since 1.1.0.0+{-# INLINE swapDupe #-}+swapDupe :: (VU.Unbox (Int, Int, w)) => VU.Vector (Int, Int, w) -> VU.Vector (Int, Int, w)+swapDupe = VU.concatMap (\(!u, !v, !w) -> VU.fromListN 2 [(u, v, w), (v, u, w)])++-- | \(O(n)\) Converts non-directed edges into directional edges. This is a convenient function for+-- making an input to `build'`.+--+-- ==== __Example__+-- `swapDupe'` duplicates each edge reversing the direction:+--+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)]+-- [(0,1),(1,0),(1,2),(2,1)]+--+-- Create a non-directed graph:+--+-- >>> let gr = Gr.build' 3 . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2)]+-- >>> gr `Gr.adj` 0+-- [1]+--+-- >>> gr `Gr.adj` 1+-- [0,2]+--+-- >>> gr `Gr.adj` 2+-- [1]+--+-- @since 1.1.0.0+{-# INLINE swapDupe' #-}+swapDupe' :: (VU.Unbox (Int, Int)) => VU.Vector (Int, Int) -> VU.Vector (Int, Int)+swapDupe' = VU.concatMap (\(!u, !v) -> VU.fromListN 2 [(u, v), (v, u)])++-- | \(O(n + m)\) Returns the strongly connected components.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 0 == 1 -> 2    3+-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 0), (1, 2)]+-- >>> Gr.scc gr+-- [[3],[0,1],[2]]+--+-- @since 1.1.0.0+{-# INLINE scc #-}+scc :: Csr w -> V.Vector (VU.Vector Int)+scc = ACISCC.sccCsr++-- | \(O(n \log n + m)\) Returns the lexicographically smallest topological ordering of the given+-- graph.+--+-- ==== Constraints+-- - The graph must be a DAG.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let n = 5+-- >>> let gr = Gr.build' n $ VU.fromList [(1, 2), (4, 0), (0, 3)]+-- >>> Gr.topSort n (gr `Gr.adj`)+-- [1,2,4,0,3]+--+-- @since 1.1.0.0+{-# INLINE topSort #-}+topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int+topSort n gr = runST $ do+  inDeg <- VUM.replicate n (0 :: Int)+  for_ [0 .. n - 1] $ \u -> do+    VU.forM_ (gr u) $ \v -> do+      VGM.modify inDeg (+ 1) v++  -- start from the vertices with zero in-degrees:+  que <- IS.new n+  inDeg' <- VU.unsafeFreeze inDeg+  VU.iforM_ inDeg' $ \v d -> do+    when (d == 0) $ do+      IS.insert que v++  buf <- B.new n+  let run = do+        IS.deleteMin que >>= \case+          Nothing -> pure ()+          Just u -> do+            B.pushBack buf u+            VU.forM_ (gr u) $ \v -> do+              nv <- subtract 1 <$> VGM.read inDeg v+              VGM.write inDeg v nv+              when (nv == 0) $ do+                IS.insert que v+            run++  run+  B.unsafeFreeze buf

@@ -0,0 +1,346 @@+{-# LANGUAGE RecordWildCards #-}++-- original implementaion:+-- <https://github.com/maspypy/library/blob/main/ds/hashmap.hpp>++-- | A dense, fast `Int` hash map with a fixed-sized `capacity` of \(n\). Most operations are+-- performed in \(O(1)\) time, but in average.+--+-- NOTE: The entries (key - value pairs) cannot be invalidated due to the internal implementation+-- (called /open addressing/).+--+-- ==== __Example__+-- Create a `HashMap` with `capacity` \(10\):+--+-- >>> import AtCoder.Extra.HashMap qualified as HM+-- >>> hm <- HM.new @_ @Int 10+--+-- `insert`, `lookup` and other functions are available in \(O(1)\) averaged time:+--+-- >>> HM.insert hm 0 100+-- >>> HM.insert hm 10 101+-- >>> HM.size hm+-- 2+--+-- >>> HM.lookup hm 0+-- Just 100+--+-- >>> HM.lookup hm 10+-- Just 101+--+-- @since 1.1.0.0+module AtCoder.Extra.HashMap+  ( -- * HashMap+    HashMap,++    -- * Constructors+    new,+    build,++    -- * Metadata+    capacity,+    size,++    -- * Lookups+    lookup,+    member,+    notMember,++    -- * Modifications++    -- ** Insertions+    insert,+    insertWith,+    exchange,++    -- ** Updates+    modify,+    modifyM,++    -- ** Reset+    clear,++    -- * Conversions++    -- ** Safe conversions+    keys,+    elems,+    assocs,++    -- ** Unsafe conversions+    unsafeKeys,+    unsafeElems,+    unsafeAssocs,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad (void, when)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Bit (Bit (..))+import Data.Bits (Bits (xor, (.&.)), (.>>.))+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import Data.Word (Word64)+import GHC.Stack (HasCallStack)+import Prelude hiding (lookup)++-- | A dense, fast `Int` hash map with a fixed-sized capacity of \(n\).+--+-- @since 1.1.0.0+data HashMap s a = HashMap+  { -- | Maximum number of elements.+    maxCapHM :: {-# UNPACK #-} !Int,+    -- | The number of elements that can be added.+    restCapHM :: !(VUM.MVector s Int),+    -- | Bit mask for powerset iteration on indexing.+    maskHM :: {-# UNPACK #-} !Int,+    -- | Original key to the hash index.+    keyHM :: !(VUM.MVector s Int),+    -- | Values to the hash index.+    valHM :: !(VUM.MVector s a),+    -- | Whether the slot is used or not.+    usedHM :: !(VUM.MVector s Bit)+  }++{-# INLINE decrementRestCapacity #-}+decrementRestCapacity :: (HasCallStack, PrimMonad m) => VUM.MVector (PrimState m) Int -> String -> m ()+decrementRestCapacity restCap funcName = do+  rest <- VGM.unsafeRead restCap 0+  let !_ = ACIA.runtimeAssert (rest > 0) $ "AtCoder.Extra.HashMap." ++ funcName ++ ": out of capacity"+  VGM.unsafeWrite restCap 0 (rest - 1)+  pure ()++-- | \(O(n)\) Creates a `HashMap` of capacity \(n\).+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (PrimMonad m, VU.Unbox a) => Int -> m (HashMap (PrimState m) a)+new n = do+  let !k0 = 1+  let !k = until (>= 2 * n) (* 2) k0+  let !maxCapHM = k `div` 2+  restCapHM <- VUM.replicate 1 maxCapHM+  let !maskHM = k - 1+  keyHM <- VUM.unsafeNew k+  valHM <- VUM.unsafeNew k+  usedHM <- VUM.replicate k $ Bit False+  pure HashMap {..}++-- | \(O(n)\) Creates a `HashMap` of capacity \(n\) with initial entries.+--+-- @since 1.1.0.0+{-# INLINE build #-}+build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (HashMap (PrimState m) a)+build n xs = do+  hm <- new n+  VU.forM_ xs $ \(!i, !x) -> do+    insert hm i x+  pure hm++-- | \(O(1)\) Returns the maximum number of elements the hash map can store.+--+-- @since 1.1.0.0+{-# INLINE capacity #-}+capacity :: HashMap s a -> Int+capacity = maxCapHM++-- | \(O(1)\) Returns the number of elements in the hash map.+--+-- @since 1.1.0.0+{-# INLINE size #-}+size :: (PrimMonad m) => HashMap (PrimState m) a -> m Int+size HashMap {..} = do+  !rest <- VUM.unsafeRead restCapHM 0+  pure $ maxCapHM - rest++-- | \(O(1)\) (Internal) Hash value calculation.+--+-- @since 1.1.0.0+{-# INLINE hash #-}+hash :: HashMap a s -> Int -> Int+hash hm x = fromIntegral $ (x3 `xor` (x3 .>>. 31)) .&. fromIntegral (maskHM hm)+  where+    fixedRandom, x1, x2, x3 :: Word64+    fixedRandom = 321896566547+    x1 = fromIntegral x + fixedRandom+    x2 = (x1 `xor` (x1 .>>. 30)) * 0xbf58476d1ce4e5b9+    x3 = (x2 `xor` (x2 .>>. 27)) * 0x94d049bb133111eb++-- | \(O(1)\) (Internal) Hashed slot search.+--+-- @since 1.1.0.0+{-# INLINE index #-}+index :: (PrimMonad m) => HashMap (PrimState m) a -> Int -> m Int+index hm@HashMap {..} k = inner (hash hm k)+  where+    inner !h = do+      Bit b <- VGM.read usedHM h+      -- already there?+      k' <- VGM.read keyHM h+      if b && k' /= k+        then inner $ (h + 1) .&. maskHM+        else pure h++-- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map+-- contains no mapping for the key.+--+-- @since 1.1.0.0+{-# INLINE lookup #-}+lookup :: (HasCallStack, VU.Unbox a, PrimMonad m) => HashMap (PrimState m) a -> Int -> m (Maybe a)+lookup hm@HashMap {..} k = do+  i <- index hm k+  Bit b <- VGM.read usedHM i+  if b+    then Just <$> VGM.read valHM i+    else pure Nothing++-- | \(O(1)\) Checks whether the hash map contains the element.+--+-- @since 1.1.0.0+{-# INLINE member #-}+member :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool+member hm@HashMap {..} k = do+  i <- index hm k+  Bit b <- VGM.read usedHM i+  -- TODO: is this key check necessary+  k' <- VGM.read keyHM i+  pure $ b && k' == k++-- | \(O(1)\) Checks whether the hash map does not contain the element.+--+-- @since 1.1.0.0+{-# INLINE notMember #-}+notMember :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool+notMember hm k = not <$> member hm k++-- | \(O(1)\) Inserts a \((k, v)\) pair.+--+-- @since 1.1.0.0+{-# INLINE insert #-}+insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m ()+insert hm k v = void $ exchange hm k v++-- | \(O(1)\) Inserts a \((k, v)\) pair. If the key exists, the function will insert the pair+-- \((k, f(v_{\mathrm{new}}, v_{\mathrm{old}}))\).+--+-- @since 1.1.0.0+{-# INLINE insertWith #-}+insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()+insertWith hm@HashMap {..} f k v = do+  i <- index hm k+  Bit b <- VGM.exchange usedHM i $ Bit True+  if b+    then do+      -- modify the existing entry+      VGM.modify valHM (f v) i+    else do+      -- insert the new \((k, v)\) pair+      decrementRestCapacity restCapHM "insertWith"+      VGM.write keyHM i k+      VGM.write valHM i v++-- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry+-- exists.+--+-- @since 1.1.0.0+{-# INLINE exchange #-}+exchange :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m (Maybe a)+exchange hm@HashMap {..} k v = do+  i <- index hm k+  Bit b <- VGM.exchange usedHM i $ Bit True+  if b+    then do+      -- overwrite the existing entry+      Just <$> VGM.exchange valHM i v+    else do+      -- insert the new (key, value) pair+      decrementRestCapacity restCapHM "exchange"+      VGM.write keyHM i k+      VGM.write valHM i v+      pure Nothing++-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists.+--+-- @since 1.1.0.0+{-# INLINE modify #-}+modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a) -> Int -> m ()+modify hm@HashMap {..} f k = do+  i <- index hm k+  Bit b <- VGM.read usedHM i+  when b $ do+    VGM.modify valHM f i++-- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists.+--+-- @since 1.1.0.0+{-# INLINE modifyM #-}+modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> m a) -> Int -> m ()+modifyM hm@HashMap {..} f k = do+  i <- index hm k+  Bit b <- VGM.read usedHM i+  when b $ do+    VGM.modifyM valHM f i++-- | \(O(n)\) Clears all the elements.+--+-- @since 1.1.0.0+{-# INLINE clear #-}+clear :: (PrimMonad m) => HashMap (PrimState m) a -> m ()+clear HashMap {..} = do+  VGM.set usedHM $ Bit False+  VUM.unsafeWrite restCapHM 0 maxCapHM++-- | \(O(n)\) Enumerates the keys in the hash map.+--+-- @since 1.1.0.0+{-# INLINE keys #-}+keys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int)+keys hm = VU.force <$> unsafeKeys hm++-- | \(O(n)\) Enumerates the elements (values) in the hash map.+--+-- @since 1.1.0.0+{-# INLINE elems #-}+elems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a)+elems hm = VU.force <$> unsafeElems hm++-- | \(O(n)\) Enumerates the key-value pairs in the hash map.+--+-- @since 1.1.0.0+{-# INLINE assocs #-}+assocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a))+assocs hm = VU.force <$> unsafeAssocs hm++-- | \(O(n)\) Enumerates the keys in the hash map.+--+-- @since 1.1.0.0+{-# INLINE unsafeKeys #-}+unsafeKeys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int)+unsafeKeys HashMap {..} = do+  used <- VU.unsafeFreeze usedHM+  keys_ <- VU.unsafeFreeze keyHM+  pure $ VU.ifilter (const . unBit . (used VG.!)) keys_++-- | \(O(n)\) Enumerates the elements (values) in the hash map.+--+-- @since 1.1.0.0+{-# INLINE unsafeElems #-}+unsafeElems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a)+unsafeElems HashMap {..} = do+  used <- VU.unsafeFreeze usedHM+  vals <- VU.unsafeFreeze valHM+  pure $ VU.ifilter (const . unBit . (used VG.!)) vals++-- | \(O(n)\) Enumerates the key-value pairs in the hash map.+--+-- @since 1.1.0.0+{-# INLINE unsafeAssocs #-}+unsafeAssocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a))+unsafeAssocs HashMap {..} = do+  used <- VU.unsafeFreeze usedHM+  keys_ <- VU.unsafeFreeze keyHM+  vals <- VU.unsafeFreeze valHM+  pure $ VU.ifilter (const . unBit . (used VG.!)) $ VU.zip keys_ vals

@@ -0,0 +1,338 @@+{-# LANGUAGE RecordWildCards #-}++-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\).+--+-- ==== __Example__+-- Create an `IntMap` with capacity \(10\):+--+-- >>> import AtCoder.Extra.IntMap qualified as IM+-- >>> im <- IM.new @_ @Int 10+--+-- `insert`, `delete`, `lookup` and other functions are available:+--+-- >>> IM.insert im 0 100+-- >>> IM.insert im 9 101+-- >>> IM.delete im 0+-- True+--+-- >>> IM.size im+-- 1+--+-- >>> IM.lookup im 9+-- Just 101+--+-- >>> IM.lookup im 1+-- Nothing+--+-- >>> IM.lookupGT im 5+-- Just (9,101)+--+-- @since 1.1.0.0+module AtCoder.Extra.IntMap+  ( -- * IntMap+    IntMap,++    -- * Constructors+    new,+    build,++    -- * Metadata+    capacity,+    size,+    null,++    -- * Lookups+    lookup,+    member,+    notMember,++    -- ** Compartive lookups+    lookupGE,+    lookupGT,+    lookupLE,+    lookupLT,++    -- ** Max/Min lookups+    lookupMin,+    lookupMax,++    -- * Modifications++    -- ** Insertions+    insert,+    insertWith,++    -- ** Updates+    modify,+    modifyM,++    -- ** Deletions+    delete,+    delete_,+    deleteMin,+    deleteMax,++    -- * Conversions+    keys,+    elems,+    assocs,+  )+where++import AtCoder.Extra.IntSet qualified as IS+import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Maybe (fromJust)+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)+import Prelude hiding (lookup, null)++-- | A dense, fast `Int` map implemented as a 64-ary tree that covers an interval \([0, n)\).+--+-- @since 1.1.0.0+data IntMap s a = IntMap+  { setIM :: !(IS.IntSet s),+    valIM :: !(VUM.MVector s a)+  }++-- | \(O(n)\) Creates an `IntMap` for an interval \([0, n)\).+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntMap (PrimState m) a)+new cap = do+  setIM <- IS.new cap+  valIM <- VUM.unsafeNew cap+  pure IntMap {..}++-- | \(O(n + m \log n)\) Creates an `IntMap` for an interval \([0, n)\) with initial values.+--+-- @since 1.1.0.0+{-# INLINE build #-}+build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (IntMap (PrimState m) a)+build cap xs = do+  im <- new cap+  VU.forM_ xs $ \(!k, !v) -> do+    insert im k v+  pure im++-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the map.+--+-- @since 1.1.0.0+{-# INLINE capacity #-}+capacity :: IntMap s a -> Int+capacity = IS.capacity . setIM++-- | \(O(1)\) Returns the number of elements in the map.+--+-- @since 1.1.0.0+{-# INLINE size #-}+size :: (PrimMonad m) => IntMap (PrimState m) a -> m Int+size = IS.size . setIM++-- | \(O(1)\) Returns whether the map is empty.+--+-- @since 1.1.0.0+{-# INLINE null #-}+null :: (PrimMonad m) => IntMap (PrimState m) a -> m Bool+null = IS.null . setIM++-- | \(O(\log n)\) Looks up the value for a key.+--+-- @since 1.1.0.0+{-# INLINE lookup #-}+lookup :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe a)+lookup im@IntMap {..} k = do+  member im k >>= \case+    True -> Just <$> VGM.read valIM k+    False -> pure Nothing++-- | \(O(\log n)\) Tests whether a key \(k\) is in the map.+--+-- @since 1.1.0.0+{-# INLINE member #-}+member :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool+member = IS.member . setIM++-- | \(O(\log n)\) Tests whether a key \(k\) is not in the map.+--+-- @since 1.1.0.0+{-# INLINE notMember #-}+notMember :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool+notMember = IS.notMember . setIM++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest key \(k\) such that \(k \ge k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupGE #-}+lookupGE :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))+lookupGE IntMap {..} k = do+  IS.lookupGE setIM k >>= \case+    Just i -> Just . (i,) <$> VGM.read valIM i+    Nothing -> pure Nothing++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest \(k\) such that \(k \gt k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupGT #-}+lookupGT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))+lookupGT is k = lookupGE is (k + 1)++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \le k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupLE #-}+lookupLE :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))+lookupLE IntMap {..} k = do+  IS.lookupLE setIM k >>= \case+    Just i -> Just . (i,) <$> VGM.read valIM i+    Nothing -> pure Nothing++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \lt k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupLT #-}+lookupLT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))+lookupLT is k = lookupLE is (k - 1)++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the minimum key \(k\).+--+-- @since 1.1.0.0+{-# INLINE lookupMin #-}+lookupMin :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))+lookupMin is = lookupGE is 0++-- | \(O(\log n)\) Looks up the \((k, v)\) pair with the maximum key \(k\).+--+-- @since 1.1.0.0+{-# INLINE lookupMax #-}+lookupMax :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))+lookupMax im = lookupLE im (IS.capacity (setIM im) - 1)++-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already+-- exists, it is overwritten.+--+-- @since 1.1.0.0+{-# INLINE insert #-}+insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> a -> m ()+insert IntMap {..} k v = do+  IS.insert setIM k+  VGM.write valIM k v++-- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already+-- exists, it overwritten with \(f(v_{\mathrm{new}}, v_{\mathrm{old}})\).+--+-- @since 1.1.0.0+{-# INLINE insertWith #-}+insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()+insertWith IntMap {..} f k v = do+  b <- IS.member setIM k+  if b+    then do+      VGM.modify valIM (f v) k+    else do+      IS.insert setIM k+      VGM.write valIM k v++-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already+-- does not exist, nothing is performed.+--+-- @since 1.1.0.0+{-# INLINE modify #-}+modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a) -> Int -> m ()+modify IntMap {..} f k = do+  b <- IS.member setIM k+  when b $ do+    VGM.modify valIM f k++-- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already+-- does not exist, nothing is performed.+--+-- @since 1.1.0.0+{-# INLINE modifyM #-}+modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> m a) -> Int -> m ()+modifyM IntMap {..} f k = do+  b <- IS.member setIM k+  when b $ do+    VGM.modifyM valIM f k++-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no+-- such key exists. Returns whether the key existed.+--+-- @since 1.1.0.0+{-# INLINE delete #-}+delete :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool+delete im = IS.delete (setIM im)++-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no+-- such key exists.+--+-- @since 1.1.0.0+{-# INLINE delete_ #-}+delete_ :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m ()+delete_ im = IS.delete_ (setIM im)++-- | \(O(\log n)\) Deletes the \((k, v)\) pair with the minimum key \(k\) in the map.+--+-- @since 1.1.0.0+{-# INLINE deleteMin #-}+deleteMin :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))+deleteMin is = do+  lookupMin is+    >>= mapM+      ( \(!key, !val) -> do+          delete_ is key+          pure (key, val)+      )++-- | \(O(\log n)\) Deletes the \((k, v)\) pair with maximum key \(k\) in the map.+--+-- @since 1.1.0.0+{-# INLINE deleteMax #-}+deleteMax :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))+deleteMax is = do+  lookupMax is+    >>= mapM+      ( \(!k, !v) -> do+          delete_ is k+          pure (k, v)+      )++-- | \(O(n \log n)\) Enumerates the keys in the map.+--+-- @since 1.1.0.0+{-# INLINE keys #-}+keys :: (PrimMonad m) => IntMap (PrimState m) a -> m (VU.Vector Int)+keys = IS.keys . setIM++-- | \(O(n \log n)\) Enumerates the elements (values) in the map.+--+-- @since 1.1.0.0+{-# INLINE elems #-}+elems :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector a)+elems im@IntMap {..} = do+  n <- IS.size setIM+  VU.unfoldrExactNM+    n+    ( \i -> do+        (!i', !x') <- fromJust <$> lookupGT im i+        pure (x', i')+    )+    (-1)++-- | \(O(n \log n)\) Enumerates the key-value pairs in the map.+--+-- @since 1.1.0.0+{-# INLINE assocs #-}+assocs :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector (Int, a))+assocs im@IntMap {..} = do+  n <- IS.size setIM+  VU.unfoldrExactNM+    n+    ( \i -> do+        (!i', !x') <- fromJust <$> lookupGT im i+        pure ((i', x'), i')+    )+    (-1)

@@ -0,0 +1,391 @@+{-# LANGUAGE RecordWildCards #-}++-- original implementation:+-- <https://github.com/maspypy/library/blob/main/ds/fastset.hpp>++-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).+--+-- ==== __Example__+-- Create an `IntSet` with capacity \(10\):+--+-- >>> import AtCoder.Extra.IntSet qualified as IS+-- >>> is <- IS.new @_ 10+--+-- `insert`, `delete` and other functions are available:+--+-- >>> IS.insert is 0+-- >>> IS.insert is 9+-- >>> IS.member is 9+-- True+--+-- >>> IS.delete is 0+-- True+--+-- >>> IS.size is+-- 1+--+-- >>> IS.member is 1+-- False+--+-- >>> IS.lookupGT is 5+-- Just 9+--+-- @since 1.1.0.0+module AtCoder.Extra.IntSet+  ( -- * IntSet+    IntSet,++    -- * Constructors+    new,+    build,++    -- * Metadata+    capacity,+    size,+    null,++    -- * Lookups+    member,+    notMember,++    -- ** Compartive lookups+    lookupGE,+    lookupGT,+    lookupLE,+    lookupLT,++    -- ** Max/Min lookups+    lookupMin,+    lookupMax,++    -- * Modifications++    -- ** Insertions+    insert,++    -- ** Deletions+    delete,+    delete_,+    deleteMin,+    deleteMax,++    -- * Conversions+    keys,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad (unless, void)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Bifunctor (bimap)+import Data.Bits+  ( Bits (clearBit, setBit, testBit),+    FiniteBits (countLeadingZeros, countTrailingZeros),+    (.<<.),+    (.>>.),+  )+import Data.Maybe (fromJust)+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)+import Prelude hiding (null)++-- | \(O(1)\) Retrieves the most significant bit.+--+-- >>> msbOf 0+-- -1+--+-- >>> msbOf maxBound+-- 62+--+-- >>> msbOf $ 4 + 2 + 1+-- 2+{-# INLINE msbOf #-}+msbOf :: Int -> Int+msbOf !x = 63 - countLeadingZeros x++-- | \(O(1)\) Retrieves the least significant bit.+--+-- >>> lsbOf 0+-- -1+--+-- >>> lsbOf maxBound+-- 0+--+-- >>> lsbOf $ 4 + 2 + 1+-- 0+{-# INLINE lsbOf #-}+lsbOf :: Int -> Int+lsbOf 0 = -1+lsbOf x = countTrailingZeros x++{-# INLINE wordSize #-}+wordSize :: Int+wordSize = 64++-- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).+--+-- @since 1.1.0.0+data IntSet s = IntSet+  { -- | Maximum number of elements.+    capacityIS :: {-# UNPACK #-} !Int,+    -- | The number of elements.+    sizeIS :: !(VUM.MVector s Int),+    -- | Segments.+    vecIS :: !(V.Vector (VUM.MVector s Int))+  }++-- | \(O(n)\) Creates an `IntSet` for the interval \([0, n)\).+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (PrimMonad m) => Int -> m (IntSet (PrimState m))+new capacityIS = do+  vecIS <-+    V.unfoldrExactNM+      (max 1 logSize)+      ( \len -> do+          let !len' = (len + wordSize - 1) `div` wordSize+          (,len') <$> VUM.replicate len' 0+      )+      capacityIS+  sizeIS <- VUM.replicate 1 (0 :: Int)+  pure IntSet {..}+  where+    (!_, !logSize) =+      until+        ((<= 1) . fst)+        (bimap ((`div` wordSize) . (+ (wordSize - 1))) (+ 1))+        (capacityIS, 0)++-- | \(O(n + m \log n)\) Creates an `IntSet` for the interval \([0, n)\) with initial values.+--+-- @since 1.1.0.0+{-# INLINE build #-}+build :: (PrimMonad m) => Int -> VU.Vector Int -> m (IntSet (PrimState m))+build n vs = do+  set <- new n+  VU.forM_ vs (insert set)+  pure set++-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the set.+--+-- @since 1.1.0.0+{-# INLINE capacity #-}+capacity :: IntSet s -> Int+capacity = capacityIS++-- | \(O(1)\) Returns the number of elements in the set.+--+-- @since 1.1.0.0+{-# INLINE size #-}+size :: (PrimMonad m) => IntSet (PrimState m) -> m Int+size = (`VUM.unsafeRead` 0) . sizeIS++-- | \(O(1)\) Returns whether the set is empty.+--+-- @since 1.1.0.0+{-# INLINE null #-}+null :: (PrimMonad m) => IntSet (PrimState m) -> m Bool+null = ((== 0) <$>) . size++-- | \(O(\log n)\) Tests whether \(k\) is in the set.+--+-- @since 1.1.0.0+{-# INLINE member #-}+member :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool+member IntSet {..} k+  | ACIA.testIndex k capacityIS = do+      let (!q, !r) = k `divMod` wordSize+      (`testBit` r) <$> VGM.unsafeRead (VG.unsafeHead vecIS) q+  | otherwise = pure False++-- | \(O(\log n)\) Tests whether \(k\) is not in the set.+--+-- @since 1.1.0.0+{-# INLINE notMember #-}+notMember :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool+notMember dis k = not <$> member dis k++-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \ge k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupGE #-}+lookupGE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)+lookupGE IntSet {..} i0+  | i0 >= capacityIS = pure Nothing+  | otherwise = inner 0 $ max 0 i0 -- REMARK: it's very important to keep @i@ non-negative.+  where+    inner h i+      | h >= V.length vecIS = pure Nothing+      -- ?+      | q == VUM.length (VG.unsafeIndex vecIS h) = pure Nothing+      | otherwise = do+          d <- (.>>. r) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q+          if d == 0+            then inner (h + 1) (q + 1)+            else+              Just+                <$> V.foldM'+                  ( \ !acc vec -> do+                      !dx <- lsbOf <$> VGM.unsafeRead vec acc+                      pure $ acc * wordSize + dx+                  )+                  (i + lsbOf d)+                  (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h))+      where+        (!q, !r) = i `divMod` wordSize++-- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \gt k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupGT #-}+lookupGT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)+lookupGT is k = lookupGE is (k + 1)++-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \le k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupLE #-}+lookupLE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)+lookupLE IntSet {..} i0+  | i0 <= -1 = pure Nothing+  | otherwise = inner 0 $ min (capacityIS - 1) i0+  where+    inner h i+      | h >= V.length vecIS = pure Nothing+      | i == -1 = pure Nothing+      | otherwise = do+          d <- (.<<. (63 - r)) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q+          if d == 0+            then inner (h + 1) (q - 1)+            else do+              Just+                <$> V.foldM'+                  ( \ !acc vec -> do+                      !dx <- msbOf <$> VGM.unsafeRead vec acc+                      pure $ acc * wordSize + dx+                  )+                  (i - countLeadingZeros d)+                  (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h))+      where+        (!q, !r) = i `divMod` wordSize++-- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \lt k_0\).+--+-- @since 1.1.0.0+{-# INLINE lookupLT #-}+lookupLT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)+lookupLT is k = lookupLE is (k - 1)++-- | \(O(\log n)\) Looks up the minimum key.+--+-- @since 1.1.0.0+{-# INLINE lookupMin #-}+lookupMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)+lookupMin is = lookupGE is 0++-- | \(O(\log n)\) Looks up the maximum key.+--+-- @since 1.1.0.0+{-# INLINE lookupMax #-}+lookupMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)+lookupMax is = lookupLE is (capacityIS is - 1)++-- | \(O(\log n)\) Inserts a key \(k\) into the set. If an entry with the same key already exists,+-- it is overwritten.+--+-- @since 1.1.0.0+{-# INLINE insert #-}+insert :: (HasCallStack, PrimMonad m) => IntSet (PrimState m) -> Int -> m ()+insert is@IntSet {..} k = do+  b <- member is k+  unless b $ do+    VUM.unsafeModify sizeIS (+ 1) 0+    V.foldM'_+      ( \i vec -> do+          let (!q, !r) = i `divMod` wordSize+          VGM.unsafeModify vec (`setBit` r) q+          pure q+      )+      k+      vecIS+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.IntSet.insert" k capacityIS++-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists. Returns+-- whether the key existed.+--+-- @since 1.1.0.0+{-# INLINE delete #-}+delete :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool+delete is@IntSet {..} k = do+  b_ <- member is k+  if b_+    then do+      VUM.unsafeModify sizeIS (subtract 1) 0+      V.foldM'_+        ( \(!b, !i) vec -> do+            let (!q, !r) = i `divMod` wordSize+            -- TODO: early return is possible+            unless b $ do+              VGM.unsafeModify vec (`clearBit` r) q+            -- `b` remembers if any other bit was on+            b' <- (/= 0) <$> VGM.unsafeRead vec q+            pure (b', q)+        )+        (False, k)+        vecIS+      pure True+    else pure False++-- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists.+--+-- @since 1.1.0.0+{-# INLINE delete_ #-}+delete_ :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m ()+delete_ is k = void $ delete is k++-- | \(O(\log n)\) Deletes the minimum key from the set. Returns `Nothing` if the set is empty.+--+-- @since 1.1.0.0+{-# INLINE deleteMin #-}+deleteMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)+deleteMin is = do+  lookupMin is+    >>= mapM+      ( \key -> do+          delete_ is key+          pure key+      )++-- | \(O(\log n)\) Deletes the maximum key from the set. Returns `Nothing` if the set is empty.+--+-- @since 1.1.0.0+{-# INLINE deleteMax #-}+deleteMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)+deleteMax is = do+  lookupMax is+    >>= mapM+      ( \key -> do+          delete_ is key+          pure key+      )++-- | \(O(n \log n)\) Enumerates the keys in the map.+--+-- @since 1.1.0.0+{-# INLINE keys #-}+keys :: (PrimMonad m) => IntSet (PrimState m) -> m (VU.Vector Int)+keys is@IntSet {sizeIS} = do+  n <- VGM.unsafeRead sizeIS 0+  VU.unfoldrExactNM+    n+    ( \i -> do+        i' <- fromJust <$> lookupGT is i+        pure (i', i')+    )+    (-1)

@@ -0,0 +1,462 @@+{-# LANGUAGE DerivingStrategies #-}++-- original implementation:+-- <https://noimi.hatenablog.com/entry/2021/05/02/195143>++-- | Dense map covering \([0, n)\) that manages non-overlapping intervals \([l, r)\) within it. Each+-- interval has an associated value \(v\). Use @onAdd@ and @onDel@ hooks to track interval state+-- changes during `buildM`, `insertM` and `deleteM` operations.+--+-- ==== Invariant+-- Each interval is operated as a whole, similar to a persistant data structure. When part of an+-- inerval is modified, the whole interval is deleted first, and the subintervals are re-inserted.+-- It's important for tracking non-linear interval information with the @onAdd@ and @onDel@ hooks+-- (callbacks).+--+-- ==== __Example__+-- Create an `IntervalMap` that covers a half-open interval \([0, n)\):+--+-- >>> import AtCoder.Extra.IntervalMap qualified as ITM+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> import Data.Vector.Unboxed.Mutable qualified as VUM+-- >>> itm <- ITM.new @_ @Int 4+--+-- It handles range set queries in amortized \(O(\log n)\) time:+--+-- >>> ITM.insert itm 0 4 0 -- 0 0 0 0+-- >>> ITM.insert itm 1 3 1 -- 0 1 1 0+-- >>> ITM.freeze itm+-- [(0,(1,0)),(1,(3,1)),(3,(4,0))]+--+-- Track interval informations with the @onAdd@ and @onDel@ hooks:+--+-- >>> import Debug.Trace (traceShow)+-- >>> itm <- ITM.new @_ @Int 4+-- >>> let onAdd l r x = print ("onAdd", l, r, x)+-- >>> let onDel l r x = print ("onDel", l, r, x)+--+-- >>> ITM.insertM itm 0 4 0 onAdd onDel -- 0 0 0 0+-- ("onAdd",0,4,0)+--+-- >>> ITM.insertM itm 1 3 1 onAdd onDel -- 0 1 1 0+-- ("onDel",0,4,0)+-- ("onAdd",0,1,0)+-- ("onAdd",3,4,0)+-- ("onAdd",1,3,1)+--+-- >>> ITM.deleteM itm 0 4 onAdd onDel+-- ("onDel",0,1,0)+-- ("onDel",1,3,1)+-- ("onDel",3,4,0)+--+-- @since 1.1.0.0+module AtCoder.Extra.IntervalMap+  ( -- * IntervalMap+    IntervalMap,++    -- * Constructors+    new,+    build,+    buildM,++    -- * Metadata+    capacity,++    -- * Lookups+    contains,+    intersects,+    lookup,+    read,+    readMaybe,++    -- * Modifications++    -- ** Insertions+    insert,+    insertM,++    -- ** Deletions+    delete,+    deleteM,++    -- ** Overwrites+    overwrite,+    overwriteM,++    -- * Conversions+    freeze,+  )+where++import AtCoder.Extra.IntMap qualified as IM+import Control.Monad (foldM_)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Vector.Generic qualified as G+import Data.Vector.Unboxed qualified as VU+import GHC.Stack (HasCallStack)+import Prelude hiding (lookup, read)++-- | Dense map covering \([0, n)\) that manages non-overlapping intervals \((l, r)\) within it. Each+-- interval has an associated value \(x\). Use @onAdd@ and @onDel@ hooks to track interval state+-- changes during `buildM`, `insertM` and `deleteM` operations.+--+-- @since 1.1.0.0+newtype IntervalMap s a = IntervalMap+  { -- | Maps \(l\) to \((r, a)\).+    unITM :: IM.IntMap s (Int, a)+  }++-- | \(O(n)\) Creates an empty `IntervalMap`.+--+-- @since 1.1.0.0+new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntervalMap (PrimState m) a)+new = fmap IntervalMap . IM.new++-- | \(O(n + m \log n)\) Creates an `IntervalMap` by combining consecutive equal values into one+-- interval.+--+-- ==== __Example__+-- >>> itm <- build @_ @Int (VU.fromList [10,10,11,11,12,12])+-- >>> freeze itm+-- [(0,(2,10)),(2,(4,11)),(4,(6,12))]+--+-- @since 1.1.0.0+build :: (PrimMonad m, Eq a, VU.Unbox a) => VU.Vector a -> m (IntervalMap (PrimState m) a)+build xs = buildM xs onAdd+  where+    onAdd _ _ _ = pure ()++-- | \(O(n + m \log n)\) Creates an `IntervalMap` by combining consecutive equal values into one+-- interval, while performing @onAdd@ hook for each interval.+--+-- @since 1.1.0.0+buildM ::+  (PrimMonad m, Eq a, VU.Unbox a) =>+  -- | Input values+  VU.Vector a ->+  -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  -- | The map+  m (IntervalMap (PrimState m) a)+buildM xs onAdd = do+  dim <- IM.new (G.length xs)+  foldM_ (step dim) (0 :: Int) $ G.group xs+  pure $ IntervalMap dim+  where+    step dim !l !xs' = do+      let !l' = l + G.length xs'+      IM.insert dim l (l', G.head xs')+      onAdd l l' (G.head xs')+      pure l'++-- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is managed by the map.+--+-- @since 1.1.0.0+{-# INLINE capacity #-}+capacity :: IntervalMap s a -> Int+capacity = IM.capacity . unITM++-- | \(O(\log n)\) Returns whether a point \(x\) is contained within any of the intervals.+--+-- @since 1.1.0.0+contains :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> m Bool+contains itm i = intersects itm i (i + 1)++-- | \(O(\log n)\) Returns whether an interval \([l, r)\) is fully contained within any of the+-- intervals.+--+-- @since 1.1.0.0+intersects :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool+intersects (IntervalMap dim) l r+  | l >= r = pure False+  | otherwise = do+      res <- IM.lookupLE dim l+      pure $ case res of+        Just (!_, (!r', !_)) -> r <= r'+        _ -> False++-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\).+--+-- @since 1.1.0.0+lookup :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe (Int, Int, a))+lookup (IntervalMap im) l r+  | l >= r = pure Nothing+  | otherwise = do+      res <- IM.lookupLE im l+      pure $ case res of+        Just (!l', (!r', !a))+          | r <= r' -> Just (l', r', a)+        _ -> Nothing++-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value.+-- Throws an error if no such interval exists.+--+-- @since 1.1.0.0+read :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m a+read itm l r = do+  res <- readMaybe itm l r+  pure $ case res of+    Just !a -> a+    Nothing -> error $ "[read] not a member: " ++ show (l, r)++-- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value.+-- Returns `Nothing` if no such interval exists.+--+-- @since 1.1.0.0+readMaybe :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe a)+readMaybe (IntervalMap dim) l r+  | l >= r = pure Nothing+  | otherwise = do+      res <- IM.lookupLE dim l+      pure $ case res of+        Just (!_, (!r', !a))+          | r <= r' -> Just a+        _ -> Nothing++-- | Amortized \(O(\log n)\) Inserts an interval \([l, r)\) with associated value \(v\) into the+-- map. Overwrites any overlapping intervals.+--+-- @since 1.1.0.0+insert :: (PrimMonad m, Eq a, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> a -> m ()+insert itm l r x = insertM itm l r x onAdd onDel+  where+    onAdd _ _ _ = pure ()+    onDel _ _ _ = pure ()++-- | Amortized \(O(\log n)\) Inserts an interval \([l, r)\) with associated value \(v\) into the+-- map. Overwrites any overlapping intervals. Tracks interval state changes via @onAdd@ and @onDel@+-- hooks.+--+-- @since 1.1.0.0+insertM ::+  (PrimMonad m, Eq a, VU.Unbox a) =>+  -- | The map+  IntervalMap (PrimState m) a ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(v\)+  a ->+  -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  m ()+insertM (IntervalMap dim) l0 r0 x onAdd onDel+  | l0 >= r0 = pure ()+  | otherwise = do+      !r <- handleRight l0 r0+      (!l', !r') <- handleLeft l0 r+      onAdd l' r' x+      IM.insert dim l' (r', x)+  where+    handleRight l r = do+      res <- IM.lookupGE dim l+      case res of+        Just interval0@(!_, (!_, !_)) -> run interval0 l r+        Nothing -> pure r++    -- Looks into intervals with @l' >= l0@.+    --           [----]+    -- (i)            *--------]   overwrite if it's x+    -- (ii)   [-------]*      delete anyways+    -- (iii)    *(------]     overwrite if it's x, or+    run (!l', (!r', !x')) l r+      | l' > r = do+          -- not adjacent: end.+          pure r+      -- (i)+      | l' == r && x' == x = do+          -- adjacent interval with the same value: merge into one.+          onDel l' r' x'+          IM.delete_ dim l'+          pure r'+      | l' == r = do+          -- adjacent interval with different values: nothing to do.+          pure r+      -- (ii)+      | r' <= r = do+          -- inside the interval: delete and continue+          onDel l' r' x'+          IM.delete_ dim l'+          res <- IM.lookupGT dim l'+          case res of+            Just rng -> run rng l r+            Nothing -> pure r+      -- (iii)+      | x' == x = do+          -- intersecting interval with the same value: merge into one.+          onDel l' r' x'+          IM.delete_ dim l'+          pure r'+      | otherwise = do+          -- intersecting interval with a different value: delete the intersection.+          onDel l' r' x'+          onAdd r r' x'+          IM.delete_ dim l'+          IM.insert dim r (r', x')+          pure r++    handleLeft l r = do+      res <- IM.lookupLT dim l+      case res of+        Nothing -> pure (l, r)+        Just (!l', (!r', !x'))+          -- (i): adjacent interval+          | r' == l && x' == x -> do+              -- adjacent interval with the same value: merge into one.+              onDel l' r' x'+              IM.delete_ dim l'+              pure (l', r)+          | r' == l -> do+              -- adjacent interval with different values: nothing to do.+              pure (l, r)+          -- (ii): not adjacent or intersecting+          | r' < l -> do+              pure (l, r)+          -- (iii): intersecting+          | x' == x -> do+              -- insersecting interval with the same value: merge into one.+              onDel l' r' x'+              IM.delete_ dim l'+              pure (min l l', max r r')+          | r' > r -> do+              -- [l', r') contains [l, r) with a different value: split into three.+              onDel l' r' x'+              onAdd l' l x'+              onAdd r r' x'+              -- IM.delete_ dim l'+              IM.insert dim l' (l, x')+              IM.insert dim r (r', x')+              pure (l, r)+          | otherwise -> do+              -- insersecting interval with a different value: delete.+              onDel l' r' x'+              onAdd l' l x'+              -- IM.delete_ dim l'+              IM.insert dim l' (l, x')+              pure (l, r)++-- | Amortized \(O(\log n)\) Deletes an interval \([l, r)\) from the map.+--+-- @since 1.1.0.0+delete :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m ()+delete itm l r = deleteM itm l r onAdd onDel+  where+    onAdd _ _ _ = pure ()+    onDel _ _ _ = pure ()++-- | Amortized \(O(\log n)\) Deletes an interval \([l, r)\) from the map. Tracks interval state+-- changes via @onAdd@ and @onDel@ hooks.+--+-- @since 1.1.0.0+deleteM ::+  (PrimMonad m, VU.Unbox a) =>+  -- | The map+  IntervalMap (PrimState m) a ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  m ()+deleteM (IntervalMap dim) l0 r0 onAdd onDel+  | l0 >= r0 = pure ()+  | otherwise = do+      handleRight l0 r0+      handleLeft l0 r0+  where+    handleRight l r = do+      res <- IM.lookupGE dim l+      case res of+        Just interval0@(!_, (!_, !_)) -> run interval0 l r+        Nothing -> pure ()++    run (!l', (!r', !x')) l r+      | l' >= r = do+          -- not intersecting+          pure ()+      | r' <= r = do+          -- contained+          onDel l' r' x'+          IM.delete_ dim l'+          res <- IM.lookupGT dim l'+          case res of+            Just rng -> run rng l r+            Nothing -> pure ()+      | otherwise = do+          -- intersecting+          onDel l' r' x'+          onAdd r r' x'+          IM.delete_ dim l'+          IM.insert dim r (r', x')+          pure ()++    handleLeft l r = do+      res <- IM.lookupLT dim l+      case res of+        Nothing -> pure ()+        Just (!l', (!r', !x'))+          | r' <= l -> do+              -- not intersecting+              pure ()+          | r' > r -> do+              -- [l', r') contains [l, r)+              onDel l' r' x'+              onAdd l' l x'+              onAdd r r' x'+              -- IM.delete dim l'+              IM.insert dim l' (l, x')+              IM.insert dim r (r', x')+          | otherwise -> do+              -- intersecting+              onDel l' r' x'+              onAdd l' l x'+              -- IM.delete_ dim l'+              IM.insert dim l' (l, x')++-- | \(O(\log n)\) Shorthand for overwriting the value of an interval that contains \([l, r)\).+--+-- @since 1.1.0.0+overwrite :: (PrimMonad m, Eq a, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> a -> m ()+overwrite itm l r x = do+  res <- lookup itm l r+  case res of+    Just (!l', !r', !_) -> insert itm l' r' x+    Nothing -> pure ()++-- | \(O(\log n)\). Shorthand for overwriting the value of an interval that contains \([l, r)\).+-- Tracks interval state changes via @onAdd@ and @onDel@ hooks.+--+-- @since 1.1.0.0+overwriteM ::+  (PrimMonad m, Eq a, VU.Unbox a) =>+  -- | The map+  IntervalMap (PrimState m) a ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(v\)+  a ->+  -- | @onAdd@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  -- | @onDel@ hook that take an interval \([l, r)\) with associated value \(v\)+  (Int -> Int -> a -> m ()) ->+  m ()+overwriteM itm l r x onAdd onDel = do+  res <- lookup itm l r+  case res of+    Just (!l', !r', !_) -> insertM itm l' r' x onAdd onDel+    Nothing -> pure ()++-- | \(O(n \log n)\) Enumerates the intervals and the associated values as \((l, (r, x))\) tuples,+-- where \([l, r)\) is the interval and \(x\) is the associated value.+--+-- @since 1.1.0.0+freeze :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> m (VU.Vector (Int, (Int, a)))+freeze = IM.assocs . unITM

@@ -1,34 +1,50 @@ -- | Extra math module. ----- ==== Examples--- >>> import AtCoder.Extra.Math qualified as M--- >>> import Data.Semigroup (Product(..), Sum(..))--- >>> getProduct $ M.power (<>) 32 (Product 2)--- 4294967296------ >>> getProduct $ M.stimes' 32 (Product 2)--- 4294967296------ >>> getProduct $ M.mtimes' 32 (Product 2)--- 4294967296--- -- @since 1.0.0.0 module AtCoder.Extra.Math-  ( -- * Binary exponential+  ( -- * Re-exports from the internal math module+    isPrime32,+    ACIM.invGcd,+    ACIM.primitiveRoot,++    -- * Binary exponentiation++    -- | ==== __Examples__+    -- >>> import AtCoder.Extra.Math qualified as M+    -- >>> import Data.Semigroup (Product(..), Sum(..))+    -- >>> getProduct $ M.power (<>) 32 (Product 2)+    -- 4294967296+    --+    -- >>> getProduct $ M.stimes' 32 (Product 2)+    -- 4294967296+    --+    -- >>> getProduct $ M.mtimes' 32 (Product 2)+    -- 4294967296     power,     stimes',     mtimes',   ) where +import AtCoder.Internal.Math qualified as ACIM import Data.Bits ((.>>.)) --- TODO: add `HasCallStack` and provide with `unsafePower`.+-- | \(O(k \log^3 n) (k = 3)\). Returns whether the given `Int` value is a prime number.+--+-- ==== Constraints+-- - \(n < 4759123141 (2^{32} < 4759123141)\), otherwise the return value can lie+--   (see [Wikipedia](https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases)).+--+--+-- @since 1.1.0.0+{-# INLINE isPrime32 #-}+isPrime32 :: Int -> Bool+isPrime32 = ACIM.isPrime --- | Calculates \(s^n\) with custom multiplication operator using the binary exponentiation+-- | Calculates \(x^n\) with custom multiplication operator using the binary exponentiation -- technique. ----- The internal implementation is taken from `Data.Semigroup.stimes`, but `power` uses strict+-- The internal implementation is taken from @Data.Semigroup.stimes@, but `power` uses strict -- evaluation and is often much faster. -- -- ==== Complexity@@ -53,7 +69,7 @@       | n == 1 = x `op` z       | otherwise = g (x `op` x) (n .>>. 1) (x `op` z) --- | Strict `Data.Semigroup.stimes`.+-- | Strict variant of @Data.Semigroup.stimes@. -- -- ==== Complexity -- - \(O(\log n)\)@@ -66,7 +82,7 @@ stimes' :: (Semigroup a) => Int -> a -> a stimes' = power (<>) --- | Strict `Data.Monoid.mtimes`.+-- | Strict variant of @Data.Monoid.mtimes@. -- -- ==== Complexity -- - \(O(\log n)\)

@@ -1,31 +1,42 @@ {-# LANGUAGE TypeFamilies #-} --- | Extra module of pre-defined `SegAct` instances.------ Be warned that they're not 100% guaranteed to be correct.+-- | Extra module of pre-defined `SegAct` instances and helpful monoids. -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid-  ( -- * SegAct (re-export)+  ( -- * Re-exports++    -- | It's mainly a list. It is recommended to use specific submodules.++    -- ** SegAct     SegAct (..), -    -- * Affine1+    -- ** Affine1     Affine1 (..),     Affine1Repr, -    -- * Range add+    -- ** Mat2x2+    Mat2x2 (..),+    Mat2x2Repr,+    V2 (..),+    V2Repr,++    -- ** Range add     RangeAdd (..),-    RangeAddId (..), -    -- * Range set+    -- ** Range set     RangeSet (..),-    RangeSetId (..),+    RangeSetRepr,++    -- ** Rolling hash+    RollingHash,   ) where  import AtCoder.Extra.Monoid.Affine1 (Affine1 (..), Affine1Repr)+import AtCoder.Extra.Monoid.Mat2x2 (Mat2x2 (..), Mat2x2Repr) import AtCoder.Extra.Monoid.RangeAdd (RangeAdd (..))-import AtCoder.Extra.Monoid.RangeAddId (RangeAddId (..))-import AtCoder.Extra.Monoid.RangeSet (RangeSet (..))-import AtCoder.Extra.Monoid.RangeSetId (RangeSetId (..))+import AtCoder.Extra.Monoid.RangeSet (RangeSet (..), RangeSetRepr)+import AtCoder.Extra.Monoid.RollingHash (RollingHash)+import AtCoder.Extra.Monoid.V2 (V2 (..), V2Repr) import AtCoder.LazySegTree (SegAct (..))

@@ -1,17 +1,23 @@ {-# LANGUAGE TypeFamilies #-} --- | Range add monoid action for \([l, r)\) intervals: \(f: x \rightarrow ax + b\).+-- | Monoid action \(f: x \rightarrow ax + b\). --+-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid+-- length in the acted monoid (@V2@).+-- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.Affine1   ( -- * Affine1     Affine1 (..),     Affine1Repr, -    -- * Constructor+    -- * Constructors     new,+    unAffine1,+    ident,+    zero, -    -- * Action+    -- * Actions     act,   ) where@@ -27,17 +33,16 @@ import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM --- Tuple is not the fastest representation, but it's easier to implement `Unbox`.---- | Range add monoid action for \([l, r)\) intervals: \(f: x \rightarrow ax + b\).+-- | Monoid action \(f: x \rightarrow ax + b\). --+-- - Use @Mat2x2@ if inverse operations are required, or if it's necessary to store the monoid+-- length in the acted monoid (@V2@).+-- -- ==== Composition and dual--- `Semigroup` for `Affine1` is implemented like function composition, and rightmost affine--- transformation is applied first: \((f_1 \circ f_2) v := f_1 (f_2(v))\). If you need 'foldr'--- of \([f_l, f_{l+1}, .., f_r)\) on a segment tree, be sure to wrap `Affine1` in--- `Data.Monoid.Dual`.+-- The affine transformation acts as a left monoid action: \(f_2 (f_1 v) = (f_2 \circ f_1) v\). To+-- apply the leftmost transformation first in a segment tree, wrap `Affine1` in @Data.Monoid.Dual@. ----- ==== Example+-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), Affine1(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> seg <- LST.build @_ @(Affine1 Int) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2]@@ -62,21 +67,42 @@ -- @since 1.0.0.0 type Affine1Repr a = (a, a) --- | Creates `Affine1`.+-- | \(O(1)\) Creates a one-dimensional affine transformation: \(f: x \rightarrow a \times x + b\). -- -- @since 1.0.0.0 {-# INLINE new #-} new :: a -> a -> Affine1 a new !a !b = Affine1 (a, b) --- | Applies \(f: x \rightarrow a \times x + b\).+-- | \(O(1)\) Retrieves the two components of `Affine1`. --+-- @since 1.1.0.0+{-# INLINE unAffine1 #-}+unAffine1 :: Affine1 a -> Affine1Repr a+unAffine1 (Affine1 a) = a++-- | \(O(1)\) Identity transformation.+--+-- @since 1.1.0.0+{-# INLINE ident #-}+ident :: (Num a) => Affine1 a+ident = Affine1 (1, 0)++-- | \(O(1)\) Transformation to zero.+--+-- @since 1.1.0.0+{-# INLINE zero #-}+zero :: (Num a) => Affine1 a+zero = Affine1 (0, 0)++-- | \(O(1)\) Applies the one-dimensional affine transformation \(f: x \rightarrow a \times x + b\).+-- -- @since 1.0.0.0 {-# INLINE act #-} act :: (Num a) => Affine1 a -> a -> a act (Affine1 (!a, !b)) x = a * x + b --- | Acts on @a@ with length in terms of `SegAct`. Works for `Sum a` only.+-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. Works for `Sum a` only. -- -- @since 1.0.0.0 {-# INLINE actWithLength #-}

@@ -0,0 +1,202 @@+{-# LANGUAGE TypeFamilies #-}++-- | Monoid action \(f: x \rightarrow ax + b\). Less efficient than @Affine1@, but compatible with+-- inverse opereations.+--+-- @since 1.1.0.0+module AtCoder.Extra.Monoid.Mat2x2+  ( -- * Mat2x2+    Mat2x2 (..),+    Mat2x2Repr,++    -- * Constructors+    new,+    unMat2x2,+    ident,+    zero,++    -- * Actions+    act,++    -- * Operators+    map,+    det,+    inv,+  )+where++import AtCoder.Extra.Math qualified as ACEM+import AtCoder.Extra.Monoid.V2 (V2 (..))+import AtCoder.LazySegTree (SegAct (..))+import Data.Semigroup (Dual (..), Semigroup (..))+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)+import Prelude hiding (map)++-- | Monoid action \(f: x \rightarrow ax + b\). Less efficient than @Affine1@, but compatible with+-- inverse opereations.+--+-- ==== Composition and dual+-- The affine transformation acts as a left monoid action: \(f_2 (f_1 v) = (f_2 \circ f_1) v\). To+-- apply the leftmost transformation first in a segment tree, wrap `Mat2x2` in @Data.Monoid.Dual@.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Monoid.Mat2x2 qualified as Mat2x2+-- >>> import AtCoder.Extra.Monoid.V2 qualified as V2+-- >>> import AtCoder.Extra.Monoid (SegAct(..), Mat2x2(..), V2(..))+-- >>> import AtCoder.LazySegTree qualified as LST+-- >>> seg <- LST.build @_ @(Mat2x2 Int) @(V2 Int) $ VU.generate 3 V2.new -- [0, 1, 2]+-- >>> LST.applyIn seg 0 3 $ Mat2x2.new 2 1 -- [1, 3, 5]+-- >>> V2.unV2 <$> LST.allProd seg+-- 9+--+-- @since 1.1.0.0+newtype Mat2x2 a = Mat2x2 (Mat2x2Repr a)+  deriving newtype+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Ord,+      -- | @since 1.1.0.0+      Show+    )++-- | `Mat2x2` internal representation. Tuples are not the fastest representation, but it's easier+-- to implement `Data.Vector.Unboxed.Unbox`.+--+-- @since 1.1.0.0+type Mat2x2Repr a = (a, a, a, a)++-- | \(O(1)\) Creates a one-dimensional affine transformation: \(f: x \rightarrow a \times x + b\).+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (Num a) => a -> a -> Mat2x2 a+new !a !b = Mat2x2 (a, b, 0, 1)++-- | \(O(1)\) Retrieves the four components of `Mat2x2`.+--+-- @since 1.1.0.0+{-# INLINE unMat2x2 #-}+unMat2x2 :: Mat2x2 a -> Mat2x2Repr a+unMat2x2 (Mat2x2 a) = a++-- | \(O(1)\) Transformation to zero.+--+-- @since 1.1.0.0+{-# INLINE zero #-}+zero :: (Num a) => Mat2x2 a+zero = Mat2x2 (0, 0, 0, 0)++-- | \(O(1)\) Identity transformation.+--+-- @since 1.1.0.0+{-# INLINE ident #-}+ident :: (Num a) => Mat2x2 a+ident = Mat2x2 (1, 0, 0, 1)++-- | \(O(1)\) Multiplies `Mat2x2` to `V2`.+{-# INLINE mulMV #-}+mulMV :: (Num a) => Mat2x2 a -> V2 a -> V2 a+mulMV (Mat2x2 (!a11, !a12, !a21, !a22)) (V2 (!x1, !x2)) = V2 (a', b')+  where+    !a' = a11 * x1 + a12 * x2+    !b' = a21 * x1 + a22 * x2++-- | \(O(1)\) Multiplies `Mat2x2` to `Mat2x2`.+{-# INLINE mulMM #-}+mulMM :: (Num a) => Mat2x2 a -> Mat2x2 a -> Mat2x2 a+mulMM (Mat2x2 (!a11, !a12, !a21, !a22)) (Mat2x2 (!b11, !b12, !b21, !b22)) = Mat2x2 (c11, c12, c21, c22)+  where+    !c11 = a11 * b11 + a12 * b21+    !c12 = a11 * b12 + a12 * b22+    !c21 = a21 * b11 + a22 * b21+    !c22 = a21 * b12 + a22 * b22++-- | \(O(1)\) Multiplies `Mat2x2` to `V2`.+--+-- @since 1.1.0.0+{-# INLINE act #-}+act :: (Num a) => Mat2x2 a -> V2 a -> V2 a+act = mulMV++-- | \(O(1)\) Maps the every component of `Mat2x2`.+--+-- @since 1.1.0.0+{-# INLINE map #-}+map :: (a -> b) -> Mat2x2 a -> Mat2x2 b+map f (Mat2x2 (!a11, !a12, !a21, !a22)) = Mat2x2 (a11', a12', a21', a22')+  where+    !a11' = f a11+    !a12' = f a12+    !a21' = f a21+    !a22' = f a22++-- | \(O(1)\) Returns the determinan of the matrix.+--+-- @since 1.1.0.0+{-# INLINE det #-}+det :: (Fractional e) => Mat2x2 e -> e+det (Mat2x2 (!a, !b, !c, !d)) = a * d - b * c++-- | \(O(1)\) Returns the inverse matrix, based on `Fractional` instance (mainly for @ModInt@).+--+-- ==== Constraints+-- - The determinant (`det`) of the matrix must be non-zero, otherwise an error is thrown.+--+-- @since 1.1.0.0+{-# INLINE inv #-}+inv :: (HasCallStack, Fractional e, Eq e) => Mat2x2 e -> Mat2x2 e+inv (Mat2x2 (!a, !b, !c, !d)) = Mat2x2 (a', b', c', d')+  where+    -- NOTE: zero division+    -- !r = recip $ a * d - b * c+    !r+      | det_ == 0 = error "AtCoder.Extra.Mat2x2.inv: the determinant of the matrix must be non zero"+      | otherwise = recip det_+      where+        !det_ = a * d - b * c+    !a' = r * d+    !b' = r * (-b)+    !c' = r * (-c)+    !d' = r * a++-- | @since 1.1.0.0+instance (Num a) => Semigroup (Mat2x2 a) where+  {-# INLINE (<>) #-}+  (<>) = mulMM+  {-# INLINE stimes #-}+  stimes = ACEM.stimes' . fromIntegral++-- | @since 1.1.0.0+instance (Num a) => Monoid (Mat2x2 a) where+  {-# INLINE mempty #-}+  mempty = ident++-- | @since 1.1.0.0+instance (Num a) => SegAct (Mat2x2 a) (V2 a) where+  {-# INLINE segAct #-}+  segAct = mulMV++-- | @since 1.1.0.0+instance (Num a) => SegAct (Dual (Mat2x2 a)) (V2 a) where+  {-# INLINE segAct #-}+  segAct (Dual f) = mulMV f++-- | @since 1.1.0.0+newtype instance VU.MVector s (Mat2x2 a) = MV_Mat2x2 (VU.MVector s (Mat2x2Repr a))++-- | @since 1.1.0.0+newtype instance VU.Vector (Mat2x2 a) = V_Mat2x2 (VU.Vector (Mat2x2Repr a))++-- | @since 1.1.0.0+deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (Mat2x2 a)++-- | @since 1.1.0.0+deriving instance (VU.Unbox a) => VG.Vector VU.Vector (Mat2x2 a)++-- | @since 1.1.0.0+instance (VU.Unbox a) => VU.Unbox (Mat2x2 a)

@@ -1,6 +1,6 @@ {-# LANGUAGE TypeFamilies #-} --- | Range add monoid action for \([l, r)\) intervals.+-- | Monoid action \(f: x \rightarrow x + d\). -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.RangeAdd@@ -9,27 +9,28 @@      -- * Constructor     new,+    unRangeAdd, -    -- * Action+    -- * Actions     act,   ) where  import AtCoder.LazySegTree (SegAct (..))-import Data.Semigroup (Sum (..))+import Data.Semigroup (stimes, Sum (..), Max(..), Min(..)) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM --- | Range add monoid action.+-- | Monoid action \(f: x \rightarrow x + d\). ----- ==== Example+-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAdd(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> import Data.Semigroup (Max(..))--- >>> seg <- LST.build @_ @(RangeAdd Int) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2]--- >>> LST.applyIn seg 0 3 $ RangeAdd 5 -- [5, 6, 7]+-- >>> seg <- LST.build @_ @(RangeAdd (Sum Int)) @(Sum Int) $ VU.generate 3 Sum -- [0, 1, 2]+-- >>> LST.applyIn seg 0 3 $ RangeAdd (Sum 5) -- [5, 6, 7] -- >>> getSum <$> LST.prod seg 0 3 -- 18 --@@ -51,36 +52,51 @@ new :: a -> RangeAdd a new = RangeAdd --- | Applies one-length range add: \(f: x \rightarrow d + x\).+-- | \(O(1)\) Retrieves the internal value of `RangeAdd`. --+-- @since 1.1.0.0+{-# INLINE unRangeAdd #-}+unRangeAdd :: RangeAdd a -> a+unRangeAdd (RangeAdd a) = a++-- | \(O(1)\) Applies one-length range add: \(f: x \rightarrow d + x\).+-- -- @since 1.0.0.0 {-# INLINE act #-}-act :: (Num a) => RangeAdd a -> a -> a-act (RangeAdd dx) x = dx + x+act :: (Semigroup a) => RangeAdd a -> a -> a+act (RangeAdd dx) x = dx <> x --- | Acts on @a@ with length in terms of `SegAct`.+-- | \(O(1)\) Acts on @a@ with length in terms of `SegAct`. -- -- @since 1.0.0.0 {-# INLINE actWithLength #-}-actWithLength :: (Num a) => Int -> RangeAdd a -> a -> a-actWithLength len (RangeAdd f) x = fromIntegral len * f + x+actWithLength :: (Semigroup a) => Int -> RangeAdd a -> a -> a+actWithLength len (RangeAdd f) x = stimes len f <> x  -- | @since 1.0.0.0-instance (Num a) => Semigroup (RangeAdd a) where+instance (Semigroup a) => Semigroup (RangeAdd a) where   {-# INLINE (<>) #-}-  (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a + b+  (RangeAdd a) <> (RangeAdd b) = RangeAdd $! a <> b --- | @since 1.0.0.0-instance (Num a) => Monoid (RangeAdd a) where+-- | @since 1.1.0.0+instance (Monoid a) => Monoid (RangeAdd a) where   {-# INLINE mempty #-}-  mempty = RangeAdd 0+  mempty = RangeAdd mempty --- | @since 1.0.0.0-instance (Num a) => SegAct (RangeAdd a) (Sum a) where+-- | @since 1.1.0.0+instance (Monoid (Sum a)) => SegAct (RangeAdd (Sum a)) (Sum a) where   {-# INLINE segActWithLength #-}-  segActWithLength len a (Sum x) = Sum $! actWithLength len a x+  segActWithLength len f x = actWithLength len f x --- not works as SegAct for Product, Min, and Max.+-- | @since 1.1.0.0+instance (Monoid (Max a)) => SegAct (RangeAdd (Max a)) (Max a) where+  {-# INLINE segActWithLength #-}+  segActWithLength len f x = actWithLength len f x++-- | @since 1.1.0.0+instance (Monoid (Min a)) => SegAct (RangeAdd (Min a)) (Min a) where+  {-# INLINE segActWithLength #-}+  segActWithLength len f x = actWithLength len f x  -- | @since 1.0.0.0 newtype instance VU.MVector s (RangeAdd a) = MV_RangeAdd (VU.MVector s a)

@@ -1,95 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Range add monoid action for \([l, r)\) intervals. Works on ideomponent monoids such as `Max`--- or `Min` only.------ @since 1.0.0.0-module AtCoder.Extra.Monoid.RangeAddId-  ( -- * RangeAddId-    RangeAddId (..),-    -- * Constructor-    new,-    -- * Action-    act,-  )-where--import AtCoder.LazySegTree (SegAct (..))-import Data.Semigroup (Max (..), Min (..))-import Data.Vector.Generic qualified as VG-import Data.Vector.Generic.Mutable qualified as VGM-import Data.Vector.Unboxed qualified as VU-import Data.Vector.Unboxed.Mutable qualified as VUM---- | Range add monoid action.------ ==== Example--- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeAddId(..))--- >>> import AtCoder.LazySegTree qualified as LST--- >>> import Data.Semigroup (Max(..))--- >>> seg <- LST.build @_ @(RangeAddId Int) @(Max Int) $ VU.generate 3 Max -- [0, 1, 2]--- >>> LST.applyIn seg 0 3 $ RangeAddId 5 -- [5, 6, 7]--- >>> getMax <$> LST.prod seg 0 3--- 7------ @since 1.0.0.0-newtype RangeAddId a = RangeAddId a-  deriving newtype-    ( -- | @since 1.0.0.0-      Eq,-      -- | @since 1.0.0.0-      Ord,-      -- | @since 1.0.0.0-      Show-    )---- | Creates `RangeAddId`.------ @since 1.0.0.0-{-# INLINE new #-}-new :: a -> RangeAddId a-new = RangeAddId---- | Applies one-length range add: \(f: x \rightarrow d + x\).------ @since 1.0.0.0-{-# INLINE act #-}-act :: (Num a) => RangeAddId a -> a -> a-act (RangeAddId f) x = f + x---- | @since 1.0.0.0-instance (Num a) => Semigroup (RangeAddId a) where-  {-# INLINE (<>) #-}-  (RangeAddId a) <> (RangeAddId b) = RangeAddId $! a + b---- | @since 1.0.0.0-instance (Num a) => Monoid (RangeAddId a) where-  {-# INLINE mempty #-}-  mempty = RangeAddId 0---- | @since 1.0.0.0-instance (Num a) => SegAct (RangeAddId a) (Max a) where-  {-# INLINE segAct #-}-  segAct f (Max x) = Max $! act f x---- | @since 1.0.0.0-instance (Num a) => SegAct (RangeAddId a) (Min a) where-  {-# INLINE segAct #-}-  segAct f (Min x) = Min $! act f x---- not works as SegAct for Sum and Product.---- | @since 1.0.0.0-newtype instance VU.MVector s (RangeAddId a) = MV_RangeAddId (VU.MVector s a)---- | @since 1.0.0.0-newtype instance VU.Vector (RangeAddId a) = V_RangeAddId (VU.Vector a)---- | @since 1.0.0.0-deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (RangeAddId a)---- | @since 1.0.0.0-deriving instance (VU.Unbox a) => VG.Vector VU.Vector (RangeAddId a)---- | @since 1.0.0.0-instance (VU.Unbox a) => VU.Unbox (RangeAddId a)

@@ -1,21 +1,22 @@ {-# LANGUAGE TypeFamilies #-} --- | Range set monoid action for \([l, r)\) intervals.+-- | Monoid action \(f: x \rightarrow a\). -- -- @since 1.0.0.0 module AtCoder.Extra.Monoid.RangeSet   ( -- * RangeSet     RangeSet (..),+    RangeSetRepr, -    -- * Constructor+    -- * Constructors     new,+    unRangeSet, -    -- * Action+    -- * Actions     act,   ) where -import AtCoder.Extra.Math qualified as ACEM import AtCoder.LazySegTree (SegAct (..)) import Data.Bit (Bit (..)) import Data.Semigroup (stimes)@@ -24,9 +25,9 @@ import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM --- | Range set monoid action.+-- | Monoid action \(f: x \rightarrow a\). ----- ==== Example+-- ==== __Example__ -- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeSet(..)) -- >>> import AtCoder.LazySegTree qualified as LST -- >>> import Data.Bit (Bit (..))@@ -51,18 +52,25 @@ -- Tuples are not the fastest representation, but it's easier to implement -- `Data.Vector.Unboxed.Unbox`. ----- @since 1.0.0.0+-- @since 1.1.0.0 type RangeSetRepr a = (Bit, a) --- | Creates a new `RangeSet` action.+-- | \(O(1)\) Creates a new `RangeSet` action. -- -- @since 1.0.0.0 {-# INLINE new #-} new :: a -> RangeSet a new = RangeSet . (Bit True,) --- | Applies one-length range set: \(f: x \rightarrow y\).+-- | \(O(1)\) Retrieves the internal representation of `RangeSet`. --+-- @since 1.1.0.0+{-# INLINE unRangeSet #-}+unRangeSet :: RangeSet a -> RangeSetRepr a+unRangeSet (RangeSet a) = a++-- | \(O(1)\) Applies one-length range set: \(f: x \rightarrow y\).+-- -- @since 1.0.0.0 {-# INLINE act #-} act :: RangeSet a -> a -> a@@ -74,7 +82,7 @@ -- @since 1.0.0.0 {-# INLINE actWithLength #-} actWithLength :: (Semigroup a) => Int -> RangeSet a -> a -> a-actWithLength len (RangeSet (Bit True, !f)) _ = ACEM.power (<>) len f+actWithLength len (RangeSet (Bit True, !f)) _ = stimes len f actWithLength _ (RangeSet (Bit False, !_)) x = x  -- | @since 1.0.0.0

@@ -1,124 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Range set monoid action for \([l, r)\) intervals. Works on ideomponent monoids such as `Max`--- or `Min` only.--------- @since 1.0.0.0-module AtCoder.Extra.Monoid.RangeSetId-  ( -- * RangeSetId-    RangeSetId (..),--    -- * Constructor-    new,--    -- * Action-    act,-  )-where--import AtCoder.LazySegTree (SegAct (..))-import Data.Bit (Bit (..))-import Data.Semigroup (Max (..), Min (..), stimes)-import Data.Vector.Generic qualified as VG-import Data.Vector.Generic.Mutable qualified as VGM-import Data.Vector.Unboxed qualified as VU-import Data.Vector.Unboxed.Mutable qualified as VUM---- | Range set monoid action.------ ==== Example--- >>> import AtCoder.Extra.Monoid (SegAct(..), RangeSetId(..))--- >>> import AtCoder.LazySegTree qualified as LST--- >>> import Data.Bit (Bit (..))--- >>> import Data.Semigroup (Max(..))--- >>> seg <- LST.build @_ @(RangeSetId (Max Int)) @(Max Int) $ VU.generate 3 (Max . (+ 10)) -- [10, 11, 12]--- >>> LST.applyIn seg 0 2 $ RangeSetId (Bit True, Max 5) -- [5, 5, 12]--- >>> getMax <$> LST.prod seg 0 3--- 12------ @since 1.0.0.0-newtype RangeSetId a = RangeSetId (RangeSetIdRepr a)-  deriving newtype-    ( -- | @since 1.0.0.0-      Eq,-      -- | @since 1.0.0.0-      Ord,-      -- | @since 1.0.0.0-      Show-    )---- | `RangeSetId` internal representation. The first value represents if it is an identity action.--- Tuples are not the fastest representation, but it's easier to implement--- `Data.Vector.Unboxed.Unbox`.------ @since 1.0.0.0-type RangeSetIdRepr a = (Bit, a)---- | Creates a new `RangeSet` action.------ @since 1.0.0.0-{-# INLINE new #-}-new :: a -> RangeSetId a-new = RangeSetId . (Bit True,)---- | Applies one-length range set: \(f: x \rightarrow y\).------ @since 1.0.0.0-{-# INLINE act #-}-act :: RangeSetId a -> a -> a-act (RangeSetId (Bit True, !f)) _ = f-act (RangeSetId (Bit False, !_)) x = x---- segActWithLength works for ideomponent monoids only.---- | @since 1.0.0.0-instance Semigroup (RangeSetId a) where-  {-# INLINE (<>) #-}-  RangeSetId (Bit False, !_) <> old = old-  new_ <> _ = new_-  {-# INLINE stimes #-}-  stimes _ x = x---- The `Monoid` constraint is just for their default value.---- | @since 1.0.0.0-instance (Monoid a) => Monoid (RangeSetId a) where-  {-# INLINE mempty #-}-  mempty = RangeSetId (Bit False, mempty)-  {-# INLINE mconcat #-}-  -- find the first non-mempty-  mconcat [] = mempty-  mconcat (RangeSetId (Bit False, !_) : as) = mconcat as-  mconcat (a : _) = a---- The target is limited to ideomponent monoids. The `Monoid` constraint is just for their default--- value.---- | @since 1.0.0.0-instance (Ord a, Bounded a) => SegAct (RangeSetId (Max a)) (Max a) where-  {-# INLINE segAct #-}-  segAct = act---- The target is limited to ideomponent monoids. The `Monoid` constraint is just for their default--- value.---- | @since 1.0.0.0-instance (Ord a, Bounded a) => SegAct (RangeSetId (Min a)) (Min a) where-  {-# INLINE segAct #-}-  segAct = act---- | @since 1.0.0.0-newtype instance VU.MVector s (RangeSetId a) = MV_RangeSetId (VU.MVector s (RangeSetIdRepr a))---- | @since 1.0.0.0-newtype instance VU.Vector (RangeSetId a) = V_RangeSetId (VU.Vector (RangeSetIdRepr a))---- | @since 1.0.0.0-deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (RangeSetId a)---- | @since 1.0.0.0-deriving instance (VU.Unbox a) => VG.Vector VU.Vector (RangeSetId a)---- | @since 1.0.0.0-instance (VU.Unbox a) => VU.Unbox (RangeSetId a)

@@ -0,0 +1,121 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE TypeFamilies #-}++-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type+-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively.+--+-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and+-- \(O(1)\) slice comparison.+--+-- @since 1.1.0.0+module AtCoder.Extra.Monoid.RollingHash+  ( -- * Rolling hash+    RollingHash (..),++    -- * Constructors+    new,+    unsafeNew,+  )+where++import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Exts (proxy#)+import GHC.TypeNats (KnownNat, natVal')++-- | Rolling hash algorithm implemented as a monoid, typically stored in a segment tree. The type+-- parameters \(b\) and \(p\) represent the B-adic base and the modulus, respectively.+--+-- Combining `RollingHash` with `SegTree` enables \(O(\log |s|)\) string slice creation and+-- \(O(1)\) slice comparison.+--+--+-- ==== __Example__+-- It's convenient to define a type alias of `RollingHash`:+--+-- >>> import AtCoder.Extra.Monoid.RollingHash qualified as RH+-- >>> import AtCoder.SegTree qualified as ST+-- >>> import Data.Char (ord)+-- >>> import Data.Semigroup (Dual (..))+-- >>> type RH = RH.RollingHash 100 998244353+--+-- Let's test whether "abcba" is a palindrome:+--+-- >>> seg <- ST.build @_ @RH . VU.map (RH.unsafeNew . ord) $ VU.fromList "abcba"+-- >>> seg' <- ST.build @_ @(Dual RH) . VU.map (Dual . RH.unsafeNew . ord) $ VU.fromList "abcba"+-- >>> hash1 <- ST.prod seg 2 5       --   cba  (left to right)+-- >>> Dual hash2 <- ST.prod seg' 0 3 -- abc    (right to lett)+-- >>> hash1 == hash2+-- True+--+-- @since 1.1.0.0+data RollingHash b p = RollingHash+  { -- | The hash value.+    hashRH :: {-# UNPACK #-} !Int,+    -- | \(b^{\mathrm{length}} \bmod p\).+    nextDigitRH :: {-# UNPACK #-} !Int+  }+  deriving+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Show+    )++-- | \(O(1)\) Creates a one-length `RollingHash` from an integer.+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p+new h = RollingHash (h `mod` fromIntegral (natVal' (proxy# @p))) (fromIntegral (natVal' (proxy# @b)))++-- | \(O(1)\) Creates a one-length `RollingHash` from an integer without taking the mod.+--+-- @since 1.1.0.0+{-# INLINE unsafeNew #-}+unsafeNew :: forall b p. (KnownNat b, KnownNat p) => Int -> RollingHash b p+unsafeNew h = RollingHash h (fromIntegral (natVal' (proxy# @b)))++-- | @since 1.1.0.0+instance (KnownNat b, KnownNat p) => Semigroup (RollingHash b p) where+  -- \| \(O(1)\)+  {-# INLINE (<>) #-}+  (RollingHash !digit1 !hash1) <> (RollingHash !digit2 !hash2) = RollingHash digit' hash'+    where+      !p = fromIntegral $ natVal' (proxy# @p)+      !digit' = digit1 * digit2 `mod` p+      !hash' = (hash1 * digit2 + hash2) `mod` p++-- | @since 1.1.0.0+instance (KnownNat b, KnownNat p) => Monoid (RollingHash b p) where+  {-# INLINE mempty #-}+  mempty = RollingHash 1 0++type RHRepr = (Int, Int)++-- | @since 1.1.0.0+instance VU.IsoUnbox (RollingHash b p) RHRepr where+  {-# INLINE toURepr #-}+  toURepr (RollingHash a b) = (a, b)+  {-# INLINE fromURepr #-}+  fromURepr (!a, !b) = RollingHash a b++-- | @since 1.1.0.0+newtype instance VU.MVector s (RollingHash b p) = MV_RH (VUM.MVector s RHRepr)++-- | @since 1.1.0.0+newtype instance VU.Vector (RollingHash b p) = V_RH (VU.Vector RHRepr)++-- | @since 1.1.0.0+deriving via (RollingHash b p `VU.As` RHRepr) instance VGM.MVector VUM.MVector (RollingHash b p)++-- | @since 1.1.0.0+deriving via (RollingHash b p `VU.As` RHRepr) instance VG.Vector VU.Vector (RollingHash b p)++-- | @since 1.1.0.0+instance VU.Unbox (RollingHash b p)

@@ -0,0 +1,81 @@+{-# LANGUAGE TypeFamilies #-}++-- | A monoid acted on by `Mat2x2`, an affine transformation target.+--+-- @since 1.1.0.0+module AtCoder.Extra.Monoid.V2+  ( -- * V2+    V2 (..),+    V2Repr,++    -- * Constructor+    new,+    unV2,+  )+where++import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM++-- | A monoid acted on by `Mat2x2`, an affine transformation target.+--+-- @since 1.1.0.0+newtype V2 a = V2 (V2Repr a)+  deriving newtype+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Ord,+      -- | @since 1.1.0.0+      Show+    )++-- | `V2` internal representation. Tuples are not the fastest representation, but it's easier+-- to implement `Data.Vector.Unboxed.Unbox`.+--+-- @since 1.1.0.0+type V2Repr a = (a, a)++-- | \(O(1)\) Creates `V2` of length \(1\).+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (Num a) => a -> V2 a+new !a = V2 (a, 1)++-- | \(O(1)\) Retrieves the value of `V2`, discarding the length information.+--+-- @since 1.1.0.0+{-# INLINE unV2 #-}+unV2 :: V2 a -> a+unV2 (V2 (!a, !_)) = a++-- | @since 1.1.0.0+instance (Num a) => Semigroup (V2 a) where+  {-# INLINE (<>) #-}+  (V2 (!a1, !a2)) <> (V2 (!b1, !b2)) = V2 (a', b')+    where+      !a' = a1 + b1+      !b' = a2 + b2++-- | @since 1.1.0.0+instance (Num a) => Monoid (V2 a) where+  {-# INLINE mempty #-}+  mempty = V2 (0, 0)++-- | @since 1.1.0.0+newtype instance VU.MVector s (V2 a) = MV_V2 (VU.MVector s (V2Repr a))++-- | @since 1.1.0.0+newtype instance VU.Vector (V2 a) = V_V2 (VU.Vector (V2Repr a))++-- | @since 1.1.0.0+deriving instance (VU.Unbox a) => VGM.MVector VUM.MVector (V2 a)++-- | @since 1.1.0.0+deriving instance (VU.Unbox a) => VG.Vector VU.Vector (V2 a)++-- | @since 1.1.0.0+instance (VU.Unbox a) => VU.Unbox (V2 a)

@@ -0,0 +1,280 @@+{-# LANGUAGE RecordWildCards #-}++-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@.  Most operations are performed+-- in \(O(1)\) time, but in average.+--+-- ==== Invariant+-- The count for each key must be non-negative. An exception is thrown if this invariant is+-- violated.+--+-- ==== Capacity limitation+-- The maximum number of distinct keys that can be inserted is fixed at \(n\), even if some keys are+-- deleted. This is due to the limitation of the internal @HashMap@.+--+-- ==== __Example__+-- Create a `MultiSet` with capacity \(4\):+--+-- >>> import AtCoder.Extra.MultiSet qualified as MS+-- >>> ms <- MS.new 4+--+-- `inc` and `dec` are the primary API:+--+-- >>> MS.inc ms 10+-- >>> MS.inc ms 10+-- >>> MS.lookup ms 10+-- Just 2+--+-- >>> MS.dec ms 10+-- >>> MS.lookup ms 10+-- Just 1+--+-- Entries with zero count are considered to be non-existing:+--+-- >>> MS.dec ms 10+-- >>> MS.member ms 10+-- False+--+-- >>> MS.lookup ms 10+-- Nothing+--+-- >>> MS.size ms+-- 0+--+-- Creating a negative count results in an exception:+--+-- >>> MS.inc ms 11+-- >>> MS.sub ms 11 2+-- *** Exception: AtCoder.Extra.Multiset.sub: the count of `11` is becoming a negative value: `-1`+-- ...+--+-- Decrementing a non-existing key does nothing and does not throw an exception:+--+-- >>> MS.dec ms 12+--+-- Misc:+--+-- >>> MS.insert ms 12 112+-- >>> MS.assocs ms+-- [(11,1),(12,112)]+--+-- @since 1.1.0.0+module AtCoder.Extra.MultiSet+  ( -- * MultiSet+    MultiSet,++    -- * Construtors+    new,++    -- * Metadata+    capacity,+    size,++    -- * Lookups+    lookup,+    member,+    notMember,++    -- * Modifications+    inc,+    dec,+    add,+    sub,+    insert,+    delete,++    -- * Conversions++    -- ** Safe conversions+    keys,+    elems,+    assocs,++    -- ** Unsafe conversions+    unsafeKeys,+    unsafeElems,+    unsafeAssocs,+  )+where++import AtCoder.Extra.HashMap qualified as HM+import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Functor ((<&>))+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)+import Prelude hiding (lookup)++-- | A fast, mutable multiset for `Int` keys backed by a @HashMap@.+--+-- @since 1.1.0.0+data MultiSet s = MultiSet+  { mapMS :: !(HM.HashMap s Int),+    cntMS :: !(VUM.MVector s Int)+  }++-- | \(O(n)\) Creates a `MultiSet` with capacity \(n\).+--+-- @since 1.1.0.0+new :: (PrimMonad m) => Int -> m (MultiSet (PrimState m))+new n = do+  mapMS <- HM.new n+  cntMS <- VUM.replicate 1 0+  pure $ MultiSet {..}++-- | \(O(1)\) Returns the maximum number of distinct keys that can be inserted into the internal+-- hash map.+--+-- @since 1.1.0.0+capacity :: MultiSet s -> Int+capacity = HM.capacity . mapMS++-- | \(O(1)\) Returns the number of distinct keys with positive counts.+--+-- @since 1.1.0.0+size :: (PrimMonad m) => MultiSet (PrimState m) -> m Int+size MultiSet {..} = do+  VGM.unsafeRead cntMS 0++-- | \(O(1)\) Looks up the count for a key.+--+-- @since 1.1.0.0+lookup :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m (Maybe Int)+lookup MultiSet {..} k = do+  HM.lookup mapMS k <&> \case+    Just i | i > 0 -> Just i+    _ -> Nothing++-- | \(O(1)\) Tests whether \(k\) is in the set.+--+-- @since 1.1.0.0+member :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool+member MultiSet {..} k = do+  HM.lookup mapMS k <&> \case+    Just i -> i > 0+    _ -> False++-- | \(O(1)\) Tests whether \(k\) is not in the set.+--+-- @since 1.1.0.0+notMember :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool+notMember ms k = not <$> member ms k++-- | \(O(1)\) Increments the count of a key.+--+-- @since 1.1.0.0+inc :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()+inc ms k = add ms k 1++-- | \(O(1)\) Decrements the count of a key.+--+-- @since 1.1.0.0+dec :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()+dec ms k = sub ms k 1++-- | \(O(1)\) Increments the count of a key \(k\) by \(c\). If the key does not exist in the set,+-- the \((k, c)\) pair is inserted. If \(v\) is negative, it falls back to `sub`.+--+-- @since 1.1.0.0+add :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()+add ms@MultiSet {..} k v = case compare v 0 of+  LT -> sub ms k (-v)+  EQ -> pure ()+  GT -> do+    HM.lookup mapMS k >>= \case+      Just n ->  do+        HM.insert mapMS k $ n + v+        when (n <= 0) $ do+          VGM.unsafeModify cntMS (+ 1) 0+      Nothing -> do+        HM.insert mapMS k v+        VGM.unsafeModify cntMS (+ 1) 0++-- | \(O(1)\) Decrements the count of a key \(k\) by \(c\). If \(c\) is negative, it falls back to+-- `add`.+--+-- @since 1.1.0.0+sub :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()+sub ms@MultiSet {..} k v = case compare v 0 of+  LT -> add ms k (-v)+  EQ -> pure ()+  GT -> do+    HM.lookup mapMS k >>= \case+      Just 0 -> pure () -- ignored+      Just n -> case compare n v of+        GT -> do+          HM.insert mapMS k (n - v)+        EQ -> do+          HM.insert mapMS k 0+          VGM.unsafeModify cntMS (subtract 1) 0+        LT -> error $ "AtCoder.Extra.Multiset.sub: the count of `" ++ show k ++ "` is becoming a negative value: `" ++ show (n - v) ++ "`"+      _ -> pure ()++-- | \(O(1)\) Inserts a key-count pair into the set. `MultiSet` is actually a count map.+--+-- @since 1.1.0.0+insert :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()+insert MultiSet {..} k v+  | v <= 0 = error $ "AtCoder.Extra.Multiset.insert: new count must be positive`" ++ show k ++ "`: `" ++ show v ++ "`"+  | otherwise = do+      HM.lookup mapMS k >>= \case+        Just n | n > 0 -> do+          HM.insert mapMS k v+        _ -> do+          HM.insert mapMS k v+          VGM.unsafeModify cntMS (+ 1) 0++-- | \(O(1)\) Deletes a key. Note that it does not undo its insertion and does not increase the+-- number of distinct keys that can be inserted into the internal hash map.+--+-- @since 1.1.0.0+delete :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()+delete MultiSet {..} k = do+  HM.lookup mapMS k >>= \case+    Just i | i > 0 -> do+      HM.insert mapMS k 0+      VGM.unsafeModify cntMS (subtract 1) 0+    _ -> pure ()++-- | \(O(n)\) Enumerates the keys in the set.+--+-- @since 1.1.0.0+{-# INLINE keys #-}+keys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)+keys ms = VU.force <$> unsafeKeys ms++-- | \(O(n)\) Enumerates the counts in the set.+--+-- @since 1.1.0.0+{-# INLINE elems #-}+elems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)+elems ms = VU.force <$> unsafeElems ms++-- | \(O(n)\) Enumerates the key-count pairs in the set.+--+-- @since 1.1.0.0+{-# INLINE assocs #-}+assocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int))+assocs ms = VU.force <$> unsafeAssocs ms++-- | \(O(n)\) Enumerates the keys in the set.+--+-- @since 1.1.0.0+{-# INLINE unsafeKeys #-}+unsafeKeys :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)+unsafeKeys = (VU.mapMaybe (\(!k, !n) -> if n > 0 then Just k else Nothing) <$>) . HM.unsafeAssocs . mapMS++-- | \(O(n)\) Enumerates the counts in the set.+--+-- @since 1.1.0.0+{-# INLINE unsafeElems #-}+unsafeElems :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector Int)+unsafeElems = (VU.filter (> 0) <$>) . HM.unsafeElems . mapMS++-- | \(O(n)\) Enumerates the key-count pairs in the set.+--+-- @since 1.1.0.0+{-# INLINE unsafeAssocs #-}+unsafeAssocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int))+unsafeAssocs = (VU.filter (\(!_, !n) -> n > 0) <$>) . HM.unsafeAssocs . mapMS

@@ -0,0 +1,306 @@+{-# LANGUAGE RecordWildCards #-}++-- original implementation:+-- <https://qiita.com/drken/items/cce6fc5c579051e64fab>++-- | A potentialized disjoint set union on a [group](https://en.wikipedia.org/wiki/Group_(mathematics\))+-- under a differential constraint system. Each vertex \(v\) is assigned a potential value \(p(v)\),+-- where representatives (`leader`) of each group have a potential of `mempty`, and other vertices have+-- potentials relative to their representative.+--+-- The group type is represented as a `Monoid` with a inverse operator, passed on `new`. This+-- approach avoids defining a separate typeclass for groups.+--+-- ==== Invariant+-- New monoids always come from the left: @new <> old@. The order is important for non-commutative+-- monoid implementations.+--+-- @since 1.1.0.0+module AtCoder.Extra.Pdsu+  ( -- * Pdsu+    Pdsu (nPdsu),++    -- * Constructors+    new,++    -- * Inspection+    leader,+    pot,+    diff,+    unsafeDiff,+    same,+    canMerge,++    -- * Merging+    merge,+    merge_,++    -- * Group information+    size,+    groups,++    -- * Reset+    clear,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++-- | A potentialized disjoint set union on a [group](https://en.wikipedia.org/wiki/Group_(mathematics\))+-- under a differential constraint system. Each vertex \(v\) is assigned a potential value \(p(v)\),+--+-- ==== __Example__+-- Create a `Pdsu` with four vertices with potential type @Sum Int@. Use `negate` as the inverse+-- operator:+--+-- >>> import AtCoder.Extra.Pdsu qualified as Pdsu+-- >>> import Data.Semigroup (Sum (..))+-- >>> dsu <- Pdsu.new @_ @(Sum Int) 4 negate+--+-- The API is similar to @Dsu@, but with differential potential values:+--+-- >>> Pdsu.merge dsu 1 0 (Sum 1)  -- p(1) - p(0) := Sum 1+-- True+--+-- >>> Pdsu.merge_ dsu 2 0 (Sum 2) -- p(2) - p(0) := Sum 2+-- >>> Pdsu.leader dsu 0+-- 0+--+-- Potential values can be retrieved with `pot`:+--+-- >>> Pdsu.pot dsu 0+-- Sum {getSum = 0}+--+-- >>> Pdsu.pot dsu 1+-- Sum {getSum = 1}+--+-- >>> Pdsu.pot dsu 2+-- Sum {getSum = 2}+--+-- Difference of potentials in the same group can be retrieved with `diff`:+--+-- >>> Pdsu.diff dsu 2 1+-- Just (Sum {getSum = 1})+--+-- >>> Pdsu.diff dsu 2 3+-- Nothing+--+-- Retrieve group information with `groups`+--+-- >>> Pdsu.groups dsu+-- [[2,1,0],[3]]+--+-- @since 1.1.0.0+data Pdsu s a = Pdsu+  { -- | The number of vertices.+    nPdsu :: {-# UNPACK #-} !Int,+    -- | Parent: non-positive, size: positive+    parentOrSizePdsu :: !(VUM.MVector s Int),+    -- | Diffierencial potential of each vertex.+    potentialPdsu :: !(VUM.MVector s a),+    invertPdsu :: !(a -> a)+  }++-- | \(O(n)\) Creates a new DSU under a differential constraint system.+--+-- @since 1.1.0.0+{-# INLINE new #-}+new ::+  forall m a.+  (PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | The number of vertices+  Int ->+  -- | The inverse operator of the monoid+  (a -> a) ->+  -- | A DSU+  m (Pdsu (PrimState m) a)+new n f = Pdsu n <$> VUM.replicate n (-1 {- size 1 -}) <*> VUM.replicate n (mempty :: a) <*> pure f++-- | \(O(\alpha(n))\) Returns the representative of the connected component that contains the+-- vertex.+--+-- @since 1.1.0.0+{-# INLINE leader #-}+leader :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int+leader Pdsu {..} v0 = inner v0+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.leader" v0 nPdsu+    inner v = do+      p <- VGM.read parentOrSizePdsu v+      if {- size? -} p < 0+        then pure v+        else do+          -- NOTE(perf): Path compression.+          -- Handle the nodes closer to the root first and move them onto just under the root+          !r <- inner p+          when (p /= r) $ do+            !pp <- VGM.read potentialPdsu p+            -- Move `v` to just under the root:+            VGM.write parentOrSizePdsu v {- root -} r+            -- INVARIANT: new coming monoids always come from the left. And we're performing+            -- reverse folding.+            VGM.modify potentialPdsu (<> pp) v+          pure r++-- | \(O(\alpha(n))\) Returns \(p(v)\), the potential value of vertex \(v\) relative to the+-- reprensetative of its group.+--+-- @since 1.1.0.0+{-# INLINE pot #-}+pot :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m a+pot dsu@Pdsu {..} v1 = do+  -- Perform path compression+  _ <- leader dsu v1+  VGM.read potentialPdsu v1+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.pot" v1 nPdsu++-- | \(O(\alpha(n))\) Returns whether the vertices \(a\) and \(b\) are in the same connected+-- component.+--+-- @since 1.1.0.0+{-# INLINE same #-}+same :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m Bool+same !dsu !v1 !v2 = (==) <$> leader dsu v1 <*> leader dsu v2++-- TODO: call it unsafeDiff++-- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\)+-- if the two vertices belong to the same group. Returns `Nothing` when the two vertices are not+-- connected.+--+-- @since 1.1.0.0+{-# INLINE diff #-}+diff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m (Maybe a)+diff !dsu !v1 !v2 = do+  b <- same dsu v1 v2+  if b+    then Just <$> unsafeDiff dsu v1 v2+    else pure Nothing++-- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\)+-- if the two vertices belong to the same group. Returns meaningless value if the two vertices are+-- not connected.+--+-- @since 1.1.0.0+{-# INLINE unsafeDiff #-}+unsafeDiff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m a+unsafeDiff !dsu !v1 !v2 = do+  p1 <- pot dsu v1+  p2 <- pot dsu v2+  pure $ p1 <> invertPdsu dsu p2++-- | \(O(\alpha(n))\) Merges \(v_1\) to \(v_2\) with differential (relative) potential+-- \(\mathrm{dp}\): \(p(v1) := \mathrm{dp} \cdot p(v2)\). Returns `True` if they're newly merged.+--+-- @since 1.1.0.0+{-# INLINE merge #-}+merge :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool+merge dsu@Pdsu {..} v10 v20 !dp0 = inner v10 v20 dp0+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 nPdsu+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 nPdsu+    inner v1 v2 !dp = do+      !r1 <- leader dsu v1+      !r2 <- leader dsu v2+      if r1 == r2+        then pure False+        else do+          -- NOTE(perf): Union by size (choose smaller one for root).+          -- Another, more proper optimization would be union by rank (depth).+          !size1 <- VGM.read potentialPdsu v1+          !size2 <- VGM.read potentialPdsu v2+          if size1 >= size2+            then do+              -- Merge `r1` onto `r2`++              -- Update the size of `r1`+              !negativeSize1 <- negate {- retrieve size -} <$> VGM.read parentOrSizePdsu r1+              !negativeSize2 <- negate {- retrieve size -} <$> VGM.read parentOrSizePdsu r2+              VGM.write parentOrSizePdsu r1 ({- size -} negativeSize1 + negativeSize2)++              -- p(v1) becomes p'(v1) under r2 after merge. p(r1) becomes p'(r1).+              --     p'(v1) = dp <> p(v2)+              --     p'(v1) = p(v1) <> 'p(r1)+              -- Therefore,+              --     p'(r1) = p^{-1}(v1) <> dp <> p(v2)+              !p1 <- VGM.read potentialPdsu v1+              !p2 <- VGM.read potentialPdsu v2+              let !pr1' = invertPdsu p1 <> dp <> p2++              -- Move `r1` to just under `r2`:+              VGM.write parentOrSizePdsu r1 {- record new root -} r2+              VGM.write potentialPdsu r1 pr1'++              pure True+            else do+              inner v2 v1 $ invertPdsu dp++-- | \(O(\alpha(n))\) `merge` with the return value discarded.+--+-- @since 1.1.0.0+{-# INLINE merge_ #-}+merge_ :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m ()+merge_ !dsu !v1 !v2 !dp = do+  _ <- merge dsu v1 v2 dp+  pure ()++-- | \(O(\alpha(n))\) Returns `True` if the two vertices belong to different groups or they belong+-- to the same group under the condition \(p(v_1) = dp \cdot p(v_2)\). It's just a convenient+-- helper function.+--+-- @since 1.1.0.0+{-# INLINE canMerge #-}+canMerge :: (HasCallStack, PrimMonad m, Semigroup a, Eq a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool+canMerge !dsu !v1 !v2 !dp = do+  b <- same dsu v1 v2+  if not b+    then pure True+    else do+      !p1 <- VGM.read (potentialPdsu dsu) v1+      !p2 <- VGM.read (potentialPdsu dsu) v2+      pure $ p1 == dp <> p2++-- | \(O(\alpha(n))\) Returns the number of vertices belonging to the same group.+--+-- @since 1.1.0.0+{-# INLINE size #-}+size :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int+size !dsu !v = (negate <$>) . VGM.read (parentOrSizePdsu dsu) =<< leader dsu v++-- | \(O(n)\) Divides the graph into connected components and returns the list of them.+--+-- @since 1.1.0.0+{-# INLINE groups #-}+groups :: (PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> m (V.Vector (VU.Vector Int))+groups dsu@Pdsu {..} = do+  groupSize <- VUM.replicate nPdsu (0 :: Int)+  leaders <- VU.generateM nPdsu $ \i -> do+    li <- leader dsu i+    VGM.modify groupSize (+ 1) li+    pure li+  result <- do+    groupSize' <- VU.unsafeFreeze groupSize+    V.mapM VUM.unsafeNew $ VU.convert groupSize'+  VU.iforM_ leaders $ \i li -> do+    i' <- subtract 1 <$> VGM.read groupSize li+    VGM.write (result VG.! li) i' i+    VGM.write groupSize li i'+  V.filter (not . VU.null) <$> V.mapM VU.unsafeFreeze result++-- | \(O(n)\) Clears the `Pdsu` to the initial state.+--+-- @since 1.1.0.0+{-# INLINE clear #-}+clear :: forall m a. (PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> m ()+clear !dsu = do+  VGM.set (potentialPdsu dsu) (mempty @a)+  VGM.set (parentOrSizePdsu dsu) (-1 {- size -})

@@ -0,0 +1,264 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE RecordWildCards #-}++-- | A simple HxW matrix backed by a vector, mainly for binary exponention.+--+-- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\).+--+-- @since 1.1.0.0+module AtCoder.Extra.Semigroup.Matrix+  ( -- * Matrix+    Matrix (..),++    -- * Constructors+    new,+    zero,+    ident,+    diag,++    -- * Mapping+    map,++    -- * Multiplications+    mulToCol,+    mul,+    mulMod,+    mulMint,++    -- * Powers+    pow,+    powMod,+    powMint,+  )+where++import AtCoder.Extra.Math qualified as ACEM+import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Internal.Barrett qualified as BT+import AtCoder.ModInt qualified as M+import Data.Foldable (for_)+import Data.Semigroup (Semigroup (..))+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Exts (proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal')+import Prelude hiding (map)++-- | A simple HxW matrix backed by a vector, mainly for binary exponention.+--+-- The matrix is a left semigroup action: \(m_2 (m_1 v) = (m_2 \circ m_1) v\).+--+--+-- @since 1.1.0.0+data Matrix a = Matrix+  { -- | @since 1.1.0.0+    hM :: {-# UNPACK #-} !Int,+    -- | @since 1.1.0.0+    wM :: {-# UNPACK #-} !Int,+    -- | @since 1.1.0.0+    vecM :: !(VU.Vector a)+  }+  deriving+    ( -- | @since 1.1.0.0+      Show,+      -- | @since 1.1.0.0+      Eq+    )++-- | Type alias of a column vector.+--+-- @since 1.1.0.0+type Col a = VU.Vector a++-- | \(O(hw)\) Creates an HxW matrix.+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (HasCallStack, VU.Unbox a) => Int -> Int -> VU.Vector a -> Matrix a+new h w vec+  | VU.length vec /= h * w = error "AtCoder.Extra.Matrix: size mismatch"+  | otherwise = Matrix h w vec++-- | \(O(n^2)\) Creates an NxN zero matrix.+--+-- @since 1.1.0.0+{-# INLINE zero #-}+zero :: (VU.Unbox a, Num a) => Int -> Matrix a+zero n = Matrix n n $ VU.replicate (n * n) 0++-- | \(O(n^2)\) Creates an NxN identity matrix.+--+-- @since 1.1.0.0+{-# INLINE ident #-}+ident :: (VU.Unbox a, Num a) => Int -> Matrix a+ident n = Matrix n n $ VU.create $ do+  vec <- VUM.replicate (n * n) 0+  for_ [0 .. n - 1] $ \i -> do+    VGM.write vec (i + n * i) 1+  pure vec++-- | \(O(n^2)\) Creates an NxN diagonal matrix.+--+-- @since 1.1.0.0+{-# INLINE diag #-}+diag :: (VU.Unbox a, Num a) => Int -> VU.Vector a -> Matrix a+diag n xs = Matrix n n $ VU.create $ do+  vec <- VUM.replicate (n * n) 0+  VU.iforM_ xs $ \i x -> do+    VGM.write vec (i + n * i) x+  pure vec++-- | \(O(n^2)\) Maps the `Matrix`.+--+-- @since 1.1.0.0+{-# INLINE map #-}+map :: (VU.Unbox a, VU.Unbox b) => (a -> b) -> Matrix a -> Matrix b+map f Matrix {..} = Matrix hM wM $ VU.map f vecM++-- | \(O(hw)\) Multiplies HxW matrix to a Hx1 column vector.+--+-- @since 1.1.0.0+{-# INLINE mulToCol #-}+mulToCol :: (Num a, VU.Unbox a) => Matrix a -> Col a -> Col a+mulToCol Matrix {..} !col = VU.convert $ V.map (VU.sum . VU.zipWith (*) col) rows+  where+    !n = VU.length col+    !_ = ACIA.runtimeAssert (n == wM) "AtCoder.Extra.Matrix.mulToCol: size mismatch"+    rows = V.unfoldrExactN hM (VU.splitAt wM) vecM++-- | \(O(h_1 K w_2)\) Multiplies H1xK matrix to a KxW2 matrix.+--+-- @since 1.1.0.0+{-# INLINE mul #-}+mul :: (Num e, VU.Unbox e) => Matrix e -> Matrix e -> Matrix e+mul !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f row col = VU.sum $ VU.imap (\iRow x -> x * VG.unsafeIndex vecB (col + iRow * w')) (VU.unsafeSlice (w * row) w vecA)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mul: matrix size mismatch"++-- | \(O(h_1 w_2 K)\) Multiplies H1xK matrix to a KxW2 matrix, taking the mod.+--+-- @since 1.1.0.0+{-# INLINE mulMod #-}+mulMod :: Int -> Matrix Int -> Matrix Int -> Matrix Int+mulMod !m !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    !bt = BT.new32 $ fromIntegral m+    f row col = VU.foldl1' addMod $ VU.imap (\iRow x -> mulMod_ x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    addMod x y = (x + y) `rem` m+    mulMod_ x y = fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMod: matrix size mismatch"++-- | \(O(h_1 w_2 K)\) `mul` specialized to `M.ModInt`.+--+-- @since 1.1.0.0+{-# INLINE mulMint #-}+mulMint :: forall a. (KnownNat a) => Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMint = mulMintImpl bt+  where+    !bt = BT.new32 $ fromIntegral (natVal' (proxy# @a))++{-# INLINE mulMintImpl #-}+mulMintImpl :: forall a. (KnownNat a) => BT.Barrett -> Matrix (M.ModInt a) -> Matrix (M.ModInt a) -> Matrix (M.ModInt a)+mulMintImpl !bt !a !b =+  Matrix h w' $+    VU.unfoldrExactN+      (h * w')+      ( \(!row, !col) ->+          let !x = f row col+           in if col + 1 >= w'+                then (x, (row + 1, 0))+                else (x, (row, col + 1))+      )+      (0, 0)+  where+    f :: Int -> Int -> M.ModInt a+    f row col = VU.sum $ VU.imap (\iRow x -> mulMod_ x (VG.unsafeIndex vecB (col + (iRow * w')))) (VU.unsafeSlice (w * row) w vecA)+    mulMod_ :: M.ModInt a -> M.ModInt a -> M.ModInt a+    mulMod_ (M.ModInt x) (M.ModInt y) = M.unsafeNew . fromIntegral $ BT.mulMod bt (fromIntegral x) (fromIntegral y)+    h = hM a+    w = wM a+    h' = hM b+    vecA = vecM a+    w' = wM b+    vecB = vecM b+    !_ = ACIA.runtimeAssert (w == h') "AtCoder.Extra.Matrix.mulMint: matrix size mismatch"++-- | \(O(w n^3)\) Calculates \(M^k\).+--+-- @since 1.1.0.0+{-# INLINE pow #-}+pow :: Int -> Matrix Int -> Matrix Int+pow k mat+  | k < 0 = error "AtCoder.Extra.Matrix.powMod: the exponential must be non-negative"+  | k == 0 = ident $ hM mat+  | otherwise = ACEM.power mul k mat+  where+    !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch"++-- | \(O(w n^3)\) Calculates \(M^k\), taking the mod.+--+-- @since 1.1.0.0+{-# INLINE powMod #-}+powMod :: Int -> Int -> Matrix Int -> Matrix Int+powMod m k mat+  | k < 0 = error "AtCoder.Extra.Matrix.powMod: the exponential must be non-negative"+  | k == 0 = ident $ hM mat+  | otherwise = ACEM.power (mulMod m) k mat+  where+    !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMod: matrix size mismatch"++-- | \(O(w n^3)\) Calculates \(M^k\), specialized to `M.ModInt`.+--+-- @since 1.1.0.0+powMint :: forall m. (KnownNat m) => Int -> Matrix (M.ModInt m) -> Matrix (M.ModInt m)+powMint k mat+  | k < 0 = error "AtCoder.Extra.Matrix.powMint: the exponential must be non-negative"+  | k == 0 = ident $ hM mat+  | otherwise = ACEM.power (mulMintImpl bt) k mat+  where+    !_ = ACIA.runtimeAssert (hM mat == wM mat) "AtCoder.Extra.Matrix.powMint: matrix size mismatch"+    !bt = BT.new32 $ fromIntegral (natVal' (proxy# @m))++-- | @since 1.1.0.0+instance (Num a, VU.Unbox a) => Semigroup (Matrix a) where+  {-# INLINE (<>) #-}+  (<>) = mul+  {-# INLINE stimes #-}+  stimes = ACEM.power (<>) . fromIntegral

@@ -0,0 +1,114 @@+-- | A permutation represented by a vector, mainly for binary exponentiation.+--+-- The permutation is a left semigroup action: \(p_2 (p_1 x) = (p_2 \circ p_1) x\).+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Semigroup.Permutation qualified as Permutation+-- >>> import Data.Semigroup (Semigroup (stimes))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let perm = Permutation.new $ VU.fromList [1, 2, 3, 0]+-- >>> Permutation.act perm 1+-- 2+--+-- >>> Permutation.act (perm <> perm) 1+-- 3+--+-- >>> Permutation.act (stimes 3 perm) 1+-- 0+--+-- @since 1.1.0.0+module AtCoder.Extra.Semigroup.Permutation+  ( -- * Permutation+    Permutation (..),++    -- * Constructors+    new,+    unsafeNew,+    ident,+    zero,++    -- * Actions+    act,++    -- * Metadata+    length,+  )+where++import AtCoder.Internal.Assert qualified as ACIA+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU+import GHC.Stack (HasCallStack)+import Prelude hiding (length)++-- | A permutation represented by a vector, mainly for binary exponentiation.+--+-- The permutation is a left semigroup action: \(p_2 (p_1 x) = (p_2 \circ p_1) x\).+--+-- @since 1.1.0.0+newtype Permutation = Permutation+  { unPermutation :: VU.Vector Int+  }+  deriving newtype+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Show+    )++-- | \(O(1)\) Creates a `Permutation`, performing boundary check on input vector.+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: (HasCallStack) => VU.Vector Int -> Permutation+new xs = Permutation xs+  where+    n = VU.length xs+    !_ = VU.foldl' (\() i -> let !_ = ACIA.runtimeAssert (-1 <= i && i < n) "AtCoder.Extra.Semigroup.Permutation.new: index boundary error" in ()) () xs++-- | \(O(1)\) Creates a `Permutation`, without performing boundary check on input vector.+--+-- @since 1.1.0.0+{-# INLINE unsafeNew #-}+unsafeNew :: (HasCallStack) => VU.Vector Int -> Permutation+unsafeNew = Permutation++-- | \(O(1)\) Creates an identity `Permutation` of length \(n\).+--+-- @since 1.1.0.0+{-# INLINE ident #-}+ident :: Int -> Permutation+ident = Permutation . (`VU.generate` id)++-- | \(O(1)\) Creates a zero `Permutation` of length \(n\). It's similar to `ident`, but filled+-- with \(-1\) and invalidates corresponding slots on composition.+--+-- @since 1.1.0.0+{-# INLINE zero #-}+zero :: Int -> Permutation+zero n = Permutation $ VU.replicate n (-1)++-- | \(O(1)\) Maps an index.+--+-- @since 1.1.0.0+{-# INLINE act #-}+act :: (HasCallStack) => Permutation -> Int -> Int+act (Permutation vec) i = case vec VG.! i of+  (-1) -> i+  i' -> i'++-- | \(O(1)\) Returns the length of the internal vector.+--+-- @since 1.1.0.0+{-# INLINE length #-}+length :: (HasCallStack) => Permutation -> Int+length = VU.length . unPermutation++-- | @since 1.1.0.0+instance Semigroup Permutation where+  {-# INLINE (<>) #-}+  Permutation r2 <> Permutation r1 = Permutation $ VU.map f r1+    where+      !_ = ACIA.runtimeAssert (VU.length r2 == VU.length r1) "AtCoder.Extra.Semigroup.Permutation.(<>): legth mismatch"+      f (-1) = -1+      f i = VG.unsafeIndex r2 i

@@ -0,0 +1,207 @@+-- | Generic tree functions.+--+-- @since 1.1.0.0+module AtCoder.Extra.Tree+  ( -- * Tree folding++    -- | These function are built around the three type parameters: \(w\), \(f\) and \(a\).+    --+    -- - \(w\): Edge weight.+    -- - \(f\): Monoid action to a vertex value. These actions are created from vertex value \(a\)+    -- and edge information @(Int, w)@.+    -- - \(a\): Monoid values stored at vertices.+    fold,+    scan,+    foldReroot,+  )+where++import Data.Functor.Identity (runIdentity)+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++{-# INLINE foldImpl #-}+foldImpl ::+  forall m w f a.+  (HasCallStack, Monad m, VU.Unbox w) =>+  (Int -> VU.Vector (Int, w)) ->+  (Int -> a) ->+  (a -> (Int, w) -> f) ->+  (f -> a -> a) ->+  Int ->+  (Int -> a -> m ()) ->+  m a+foldImpl tree valAt toF act root memo = inner (-1) root+  where+    inner :: Int -> Int -> m a+    inner !parent !v1 = do+      let !acc0 = valAt v1+      let !v2s = VU.filter ((/= parent) . fst) $ tree v1+      !res <- VU.foldM' (\acc (!v2, !w) -> (`act` acc) . (`toF` (v1, w)) <$> inner v1 v2) acc0 v2s+      memo v1 res+      pure res++-- | \(O(n)\) Folds a tree from a root vertex, also known as tree DP.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree qualified as Tree+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let gr = Gr.build @(Sum Int) 5 . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]+-- >>> type W = Sum Int -- edge weight+-- >>> type F = Sum Int -- action type+-- >>> type X = Sum Int -- vertex value+-- >>> :{+--  let res = Tree.fold (gr `Gr.adjW`) valAt toF act 2+--        where+--          valAt :: Int -> X+--          valAt = const $ mempty @(Sum Int)+--          toF :: X -> (Int, W) -> F+--          toF x (!_i, !dx) = x + dx+--          act :: F -> X -> X+--          act dx x = dx + x+--   in getSum res+-- :}+-- 4+--+-- @since 1.1.0.0+{-# INLINE fold #-}+fold ::+  (HasCallStack, VU.Unbox w) =>+  -- | Graph as a function.+  (Int -> VU.Vector (Int, w)) ->+  -- | @valAt@: Assignment of initial vertex values.+  (Int -> a) ->+  -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.+  (a -> (Int, w) -> f) ->+  -- | @act@: Performs an action onto a vertex value.+  (f -> a -> a) ->+  -- | Root vertex.+  Int ->+  -- | Tree folding result from the root vertex.+  a+fold tree valAt toF act root = runIdentity $ do+  foldImpl tree valAt toF act root (\_ _ -> pure ())++-- | \(O(n)\) Folds a tree from a root vertex, also known as tree DP. The calculation process on+-- every vertex is recoreded and returned as a vector.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree qualified as Tree+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let n = 5+-- >>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]+-- >>> type W = Sum Int -- edge weight+-- >>> type F = Sum Int -- action type+-- >>> type X = Sum Int -- vertex value+-- >>> :{+--  let res = Tree.scan n (gr `Gr.adjW`) valAt toF act 2+--        where+--          valAt :: Int -> X+--          valAt = const $ mempty @(Sum Int)+--          toF :: X -> (Int, W) -> F+--          toF x (!_i, !dx) = x + dx+--          act :: F -> X -> X+--          act dx x = dx + x+--   in VU.map getSum res+-- :}+-- [0,1,4,1,0]+--+-- @since 1.1.0.0+{-# INLINE scan #-}+scan ::+  (VU.Unbox w, VG.Vector v a) =>+  -- | The number of vertices.+  Int ->+  -- | Graph as a function.+  (Int -> VU.Vector (Int, w)) ->+  -- | @valAt@: Assignment of initial vertex values.+  (Int -> a) ->+  -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.+  (a -> (Int, w) -> f) ->+  -- | @act@: Performs an action onto a vertex value.+  (f -> a -> a) ->+  -- | Root vertex.+  Int ->+  -- | Tree scanning result from a root vertex.+  v a+scan n tree acc0At toF act root = VG.create $ do+  dp <- VGM.unsafeNew n+  !_ <- foldImpl tree acc0At toF act root $ \v a -> do+    VGM.unsafeWrite dp v a+  pure dp++-- | \(O(n)\) Folds a tree from every vertex, using the rerooting technique.+--+-- ==== __Example__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree qualified as Tree+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let n = 5+-- >>> let gr = Gr.build @(Sum Int) n . Gr.swapDupe $ VU.fromList [(2, 1, Sum 1), (1, 0, Sum 1), (2, 3, Sum 1), (3, 4, Sum 1)]+-- >>> type W = Sum Int -- edge weight+-- >>> type F = Sum Int -- action type+-- >>> type X = Sum Int -- vertex value+-- >>> :{+--  let res = Tree.foldReroot n (gr `Gr.adjW`) valAt toF act+--        where+--          valAt :: Int -> X+--          valAt = const $ mempty @(Sum Int)+--          toF :: X -> (Int, W) -> F+--          toF x (!_i, !dx) = x + dx+--          act :: F -> X -> X+--          act dx x = dx + x+--   in VU.map getSum res+-- :}+-- [4,4,4,4,4]+--+-- @since 1.1.0.0+{-# INLINE foldReroot #-}+foldReroot ::+  forall w f a.+  (HasCallStack, VU.Unbox w, VU.Unbox a, VU.Unbox f, Monoid f) =>+  -- | The number of vertices.+  Int ->+  -- | Graph as a function.+  (Int -> VU.Vector (Int, w)) ->+  -- | @valAt@:Assignment of initial vertex values.+  (Int -> a) ->+  -- | @toF@: Converts a vertex value into an action onto a neighbor vertex.+  (a -> (Int, w) -> f) ->+  -- | @act@: Performs an action onto a vertex value.+  (f -> a -> a) ->+  -- | Tree folding result from every vertex as a root.+  VU.Vector a+foldReroot n tree valAt toF act = VU.create $ do+  -- Calculate tree DP for every vertex as a root:+  !dp <- VUM.unsafeNew n+  let reroot parent parentF v1 = do+        -- TODO: when the operator is not commutative?+        let !children = VU.filter ((/= parent) . fst) $ tree v1+        let !fL = VU.scanl' (\ !f (!v2, !w) -> (f <>) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children+        let !fR = VU.scanr' (\(!v2, !w) !f -> (<> f) . (`toF` (v1, w)) $ treeDp VG.! v2) f0 children++        -- save+        let !x1 = (parentF <> VU.last fL) `act` valAt v1+        VGM.unsafeWrite dp v1 x1++        VU.iforM_ children $ \i2 (!v2, !w) -> do+          -- composited operator excluding @v2@:+          let !f1 = parentF <> (fL VG.! i2) <> (fR VG.! (i2 + 1))+          let !v1Acc = f1 `act` valAt v1+          let !f2 = toF v1Acc (v2, w)+          reroot v1 f2 v2++  reroot (-1 :: Int) f0 root0+  pure dp+  where+    !root0 = 0 :: Int+    !f0 = mempty @f+    !treeDp = scan n tree valAt toF act root0 :: VU.Vector a

@@ -0,0 +1,512 @@+{-# LANGUAGE RecordWildCards #-}++-- | Heavy-light decomposition is a method for partitioning a tree into segments with consecutive+-- indices. It processes various path queries in \(O(\log n)\) time. For segment tree integration+-- and monoid products, refer to the @TreeMonoid@ module.+--+-- ==== __Overview of the internals__+-- The following is for understanding the internals, not for using the API. Skip to the examples if+-- you want.+--+-- ===== Original tree+--+-- Consider a tree with arbitrary vertex order:+--+-- @+--  0--8--7--3--1--2--12--13--15--14     XX: original Vertex+--     |        |                        --: edge+-- 10--5        11--8--6                 |: edge+--     |+--     4+-- @+--+-- ===== `indexHld`: Vertex -> VertexHld+--+-- The tree vertices are reindexed with `indexHld`, where each segment is assigned consecutive+-- vertex indices:+--+-- @+--  0==1==2==3==4==5==6==7==8==9     XX: VertexHld+--     |        |                    ==: edges on the same semgent+-- 14==13       10==11==12           |: edge between different segments+--     |+--     15+-- @+--+-- Note that vertices on higher (closer to the root) segments are assigned smaller indices. This is+-- very internally very important when calculating `lca`.+--+-- ===== `headHld`: Vertex -> Vertex+--+-- `headHld` points the "head" vertex of each segment. It can be used for finding LCA of two+-- vertices. To find the LCA, move up to the head, go up to the parental segment's vertex and+-- repeat until the two vertices are on the same segment.+--+-- @+--  0==0==0==0==0==0==0==0==0==0     XX: original Vertex+--     |     |+--  5==5     11==11==11+--     |+--     4+-- @+--+-- `headHld` also works for identifying segments. When two vertices are on the same segment, they+-- have the same head.+--+-- ===== `parentHld`: Vertex -> Vertex+--+-- `parentHld` points the parental segment's vertex from a head:+--+-- @+-- (-1)==0==8==7==3==1==2==12==13==15     XX: original Vertex+--       |        |+--    5==8        1==11=8+--       |+--       5+-- @+--+-- ==== __Example__+-- Create an `Hld` for a tree:+--+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 0--1--2--3+-- >>> --    ++-- >>> --    +--4--5+-- >>> let n = 6+-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)]+-- >>> let hld = Hld.new tree+--+-- `Hld` can process various queries in \(O(\log n)\) time:+--+-- >>> Hld.ancestor hld 5 3 -- go up three parents from `5`+-- 0+--+-- >>> Hld.jump hld 5 2 3   -- go to the third vertex from `5` to `2`:+-- Just 2+--+-- >>> Hld.lengthBetween hld 5 3 -- get the length (the number of edges) between `5` and `3`:+-- 4+--+-- >>> Hld.path hld 5 3     -- get the path between `5` and `3`:+-- [5,4,1,2,3]+--+-- Our `Hld` is rooted at @0@ vertex and subtree queries are available:+--+-- >>> Hld.isInSubtree hld 2 3 -- `3` is in the subtree of `2`+-- True+--+-- >>> Hld.isInSubtree hld 2 4 -- `4` is not in the subtree of `2`+-- False+--+-- ===== Segment queries+-- Products and segment queries are primarily used by the @TreeMonoid@ module and is not intended+-- for diretct use, but here's some examples. This time the reindex by the HLD is identity:+--+-- >>> Hld.indexHld hld+-- [0,1,2,3,4,5]+--+-- So we can easily understand the outputs:+--+-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnVertices hld 5 3+-- [(5,4),(1,3)]+--+-- >>> Hld.pathSegmentsInclusive Hld.WeightsAreOnEdges hld 5 3 -- LCA (1) is removed+-- [(5,4),(2,3)]+--+-- >>> Hld.subtreeSegmentInclusive hld 1+-- (1,5)+--+-- @since 1.1.0.0+module AtCoder.Extra.Tree.Hld+  ( -- * Hld+    Hld (..),+    Vertex,+    VertexHld,++    -- * Constructors+    new,+    newAt,++    -- * LCA+    lca,++    -- * Jump+    ancestor,+    jump,++    -- * Path+    lengthBetween,+    path,+    pathSegmentsInclusive,++    -- * Subtree+    subtreeSegmentInclusive,+    isInSubtree,++    -- * Products+    WeightPolicy (..),+    prod,+  )+where++import AtCoder.Extra.Graph qualified as Gr+import AtCoder.Internal.Assert qualified as ACIA+import Control.Monad+import Control.Monad.Fix+import Control.Monad.ST+import Data.Maybe+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++-- | Original graph vertex.+--+-- @since 1.1.0.0+type Vertex = Int++-- | Vertex reindexed by `indexHld`.+--+-- @since 1.1.0.0+type VertexHld = Vertex++-- | `Hld` partitions a tree into segments and assignes contiguous `VertexHld` for each segment.+--+-- @since 1.1.0.0+data Hld = Hld+  { -- | The root vertex.+    --+    -- @since 1.1.0.0+    rootHld :: {-# UNPACK #-} !Vertex,+    -- | Maps `Vertex` to the parent `Vertex`. Returns @-1@ for the root node.+    --+    -- @since 1.1.0.0+    parentHld :: !(VU.Vector Vertex),+    -- | Maps `Vertex` to `VertexHld`, re-indexed vertices contiguous in each segment.+    --+    -- @since 1.1.0.0+    indexHld :: !(VU.Vector VertexHld),+    -- | Maps `Vertex` to the head `Vertex` of the segment.+    --+    -- @since 1.1.0.0+    headHld :: !(VU.Vector Vertex),+    -- | Maps `VertexHld` back to `Vertex`. Used for `ancestor` etc.+    --+    -- @since 1.1.0.0+    revIndexHld :: !(VU.Vector Vertex),+    -- | Maps `Vertex` to their depth from the root. Used for `jump` etc.+    --+    -- @since 1.1.0.0+    depthHld :: !(VU.Vector Int),+    -- | Maps `Vertex` to the subtree size. This is for subtree products.+    --+    -- @since 1.1.0.0+    subtreeSizeHld :: !(VU.Vector Int)+  }+  deriving+    ( -- | @since 1.1.0.0+      Show,+      -- | @since 1.1.0.0+      Eq+    )++-- | \(O(n)\) Creates an `Hld` with \(0\) as the root vertex.+--+-- @since 1.1.0.0+{-# INLINE new #-}+new :: forall w. (HasCallStack) => Gr.Csr w -> Hld+new tree = newAt tree 0++-- | \(O(n)\) Creates an `Hld` with a root vertex specified.+--+-- @since 1.1.0.0+{-# INLINE newAt #-}+newAt :: forall w. (HasCallStack) => Gr.Csr w -> Vertex -> Hld+newAt tree root = runST $ do+  -- Re-create adjacent vertices so that the biggest subtree's head vertex comes first.+  --+  -- We /could/ instead record the biggest adjacent subtree vertex for each vertex, but the other+  -- DFS would be harder.+  let (!tree', !parent, !depths, !subtreeSize) = runST $ do+        adjVec <- VU.thaw (Gr.adjCsr tree)+        parent_ <- VUM.unsafeNew n+        depths_ <- VUM.unsafeNew n+        subtreeSize_ <- VUM.unsafeNew n++        _ <- (\f -> fix f 0 (-1) root) $ \loop depth p v1 -> do+          VGM.write parent_ v1 p+          VGM.write depths_ v1 depth++          (!size1, !eBig) <-+            VU.foldM'+              ( \(!size1, !eBig) (!e2, !v2) -> do+                  if v2 == p+                    then pure (size1, eBig)+                    else do+                      size2 <- loop (depth + 1) v1 v2+                      -- NOTE: It's `>` because we should swap at least once if there's some vertex other+                      -- that the parent_.+                      pure (size1 + size2, if size1 > size2 then eBig else e2)+              )+              (1 :: Int, -1)+              (tree `Gr.eAdj` v1)++          -- move the biggest subtree's head to the first adjacent vertex.+          -- it means the "heavy edge" or the longest segment.+          when (eBig /= -1) $ do+            VGM.swap adjVec eBig $ fst (VG.head (tree `Gr.eAdj` v1))++          -- record subtree size+          VGM.write subtreeSize_ v1 size1++          pure size1++        !vec <- VU.unsafeFreeze adjVec+        (tree {Gr.adjCsr = vec},,,)+          <$> VU.unsafeFreeze parent_+          <*> VU.unsafeFreeze depths_+          <*> VU.unsafeFreeze subtreeSize_++  -- vertex -> reindexed vertex index+  indices <- VUM.replicate n (-1 :: Int)++  -- vertex -> head vertex of the segment+  heads <- VUM.replicate n (-1 :: Int)++  _ <- (\f -> fix f (0 :: Int) root (-1) root) $ \loop acc h p v1 -> do+    -- reindex:+    VGM.write indices v1 acc+    let !acc' = acc + 1++    VGM.write heads v1 h++    -- when the first vertex is within the same segment:+    let (!adj1, !rest) = fromJust $ VU.uncons (tree' `Gr.adj` v1)+    acc'' <-+      if adj1 == p+        then pure acc'+        else loop acc' h v1 adj1++    -- the others are in other segments:+    VU.foldM'+      ( \a v2 -> do+          if v2 == p+            then pure a+            else loop a v2 v1 v2+      )+      acc''+      rest++  !indices' <- VU.unsafeFreeze indices+  let !revIndex = VU.update (VU.replicate n (-1)) $ VU.imap (flip (,)) indices'++  Hld root parent indices'+    <$> VU.unsafeFreeze heads+    <*> pure revIndex+    <*> pure depths+    <*> pure subtreeSize+  where+    !n = Gr.nCsr tree+    !_ = ACIA.runtimeAssert (2 * (Gr.nCsr tree - 1) == Gr.mCsr tree) "AtCoder.Extra.Hld.newAt: not a non-directed tree"++-- | \(O(\log n)\) Calculates the lowest common ancestor of \(u\) and \(v\).+--+-- @since 1.1.0.0+{-# INLINE lca #-}+lca :: (HasCallStack) => Hld -> Vertex -> Vertex -> Vertex+lca Hld {..} = inner+  where+    inner !x !y+      -- sort for easier processing+      -- TODO: @case compare ix iy@ would be easier for me to understand+      | ix > iy = inner y x+      -- @x@ and @y@ are in other segments:+      | hx /= hy = inner x $ parentHld VG.! hy+      -- @x@ and @y@ are within the same segment:+      -- select the smaller one, which is closer to the root and that is the LCA.+      | otherwise = x+      where+        !ix = indexHld VG.! x+        !iy = indexHld VG.! y+        hx = headHld VG.! x+        hy = headHld VG.! y++-- | \(O(\log n)\) Go up \(k\) times from a vertex \(v\) to the root node. Throws an error if \(k\)+-- is bigger than the depth of \(v\).+--+-- @since 1.1.0.0+{-# INLINE ancestor #-}+ancestor :: (HasCallStack) => Hld -> Vertex -> Int -> Vertex+ancestor Hld {..} parent k0 = inner parent k0+  where+    !_ = ACIA.runtimeAssert (0 <= k0 && k0 <= depthHld VG.! parent) $ "AtCoder.Extra.Tree.Hld.ancestor: k-th ancestor is out of the bounds (`k = " ++ show k0 ++ "`)"+    inner v k+      -- on this segment+      | k <= iv - ihv = revIndexHld VG.! (iv - k)+      -- next segment+      | otherwise = inner (parentHld VG.! hv) (k - (iv - ihv + 1))+      where+        iv = indexHld VG.! v+        hv = headHld VG.! v+        ihv = indexHld VG.! hv++-- | \(O(\log n)\) Returns the \(k\)-th vertex of the path between \(u\) and \(v\) from \(u\).+-- Throws an error if `k` is out+--+-- @since 1.1.0.0+{-# INLINE jump #-}+jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex+jump hld@Hld {..} u v k+  | k > lenU + lenV = Nothing+  | k <= lenU = Just $ ancestor hld u k+  | otherwise = Just $ ancestor hld v (lenU + lenV - k)+  where+    lca_ = lca hld u v+    du = depthHld VG.! u+    dv = depthHld VG.! v+    lenU = du - depthHld VG.! lca_+    lenV = dv - depthHld VG.! lca_++-- | \(O(\log n)\) Returns the length of the path between \(u\) and \(v\).+--+-- @since 1.1.0.0+{-# INLINE lengthBetween #-}+lengthBetween :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int+lengthBetween hld@Hld {..} u v = du - dLca + dv - dLca+  where+    !lca_ = lca hld u v+    !dLca = depthHld VG.! lca_+    !du = depthHld VG.! u+    !dv = depthHld VG.! v++-- | \(O(n)\) Returns the vertices on the path between \(u\) and \(v\).+--+-- @since 1.1.0.0+{-# INLINE path #-}+path :: (HasCallStack) => Hld -> Vertex -> Vertex -> [Vertex]+path hld@Hld {..} u v = concatMap expand $ pathSegmentsInclusive WeightsAreOnVertices hld u v+  where+    expand (!l, !r)+      | l <= r = map (revIndexHld VG.!) [l .. r]+      | otherwise = map (revIndexHld VG.!) [l, l - 1 .. r]++-- | \(O(\log n)\) Decomposes a path between two vertices \(u\) and \(v\) into segments. Each+-- segment is represented as an __inclusive__ range \([u_i, v_i]\) of `VertexHLD`.+--+-- The LCA is omitted from the returning vertices when the weight policy is set to+-- `WeightsAreOnEdges`. This is the trick to put edge weights to on vertices.+--+-- @since 1.1.0.0+{-# INLINE pathSegmentsInclusive #-}+pathSegmentsInclusive :: (HasCallStack) => WeightPolicy -> Hld -> Vertex -> Vertex -> [(VertexHld, VertexHld)]+pathSegmentsInclusive weightPolicy Hld {..} x0 y0 = done $ inner x0 [] y0 []+  where+    isEdge = weightPolicy == WeightsAreOnEdges+    done (!up, !down) = reverse up ++ down+    -- @up@: bottom to top. [(max, min)]+    -- @down@: top to bottom. [(min, max)]+    inner :: Vertex -> [(VertexHld, VertexHld)] -> Vertex -> [(VertexHld, VertexHld)] -> ([(VertexHld, VertexHld)], [(VertexHld, VertexHld)])+    inner x up y down+      | hx == hy && isEdge = case compare ix iy of+          -- skip LCA on edge vertices+          LT -> (up, (ix {- edge -} + 1, iy) : down)+          GT -> ((ix, iy {- edge -} + 1) : up, down)+          EQ -> (up, down)+      | hx == hy && not isEdge = case compare ix iy of+          LT -> (up, (ix, iy) : down)+          _ -> ((ix, iy) : up, down)+      | otherwise = case compare ix iy of+          LT -> inner x up phy ((ihy, iy) : down)+          GT -> inner phx ((ix, ihx) : up) y down+          EQ -> error "unreachable"+      where+        ix, iy :: VertexHld+        !ix = indexHld VG.! x+        !iy = indexHld VG.! y+        hx, hy :: Vertex+        hx = headHld VG.! x+        hy = headHld VG.! y+        ihx, ihy :: VertexHld+        ihx = indexHld VG.! hx+        ihy = indexHld VG.! hy+        phx, phy :: VertexHld+        phx = parentHld VG.! hx+        phy = parentHld VG.! hy++-- | \(O(1)\) Returns a half-open interval of `VertexHld` \([\mathrm{start}, \mathrm{end})\) that+-- corresponds to the subtree segments rooted at the given @subtreeRoot@.+--+-- @since 1.1.0.0+{-# INLINE subtreeSegmentInclusive #-}+subtreeSegmentInclusive :: (HasCallStack) => Hld -> Vertex -> (VertexHld, VertexHld)+subtreeSegmentInclusive Hld {..} subtreeRoot = (ir, ir + sr - 1)+  where+    ir = indexHld VG.! subtreeRoot+    sr = subtreeSizeHld VG.! subtreeRoot++-- | \(O(1)\) Returns `True` if \(u\) is in a subtree of \(r\).+--+-- @since 1.1.0.0+{-# INLINE isInSubtree #-}+isInSubtree :: (HasCallStack) => Hld -> Vertex -> Vertex -> Bool+isInSubtree hld@Hld {..} r_ u = l <= iu && iu <= r+  where+    (!l, !r) = subtreeSegmentInclusive hld r_+    !iu = indexHld VG.! u++-- | Represents whether weights are put on vertices or edges.+--+-- @since 1.1.0.0+data WeightPolicy+  = -- | Weights are put on vertices.+    --+    -- @since 1.1.0.0+    WeightsAreOnVertices+  | -- | Weights are put on edges.+    --+    -- @since 1.1.0.0+    WeightsAreOnEdges+  deriving+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Show+    )++-- | \(O(\log n f)\) Returns product of the path between \(u\) and \(v\), using the user functions+-- of time complexity \(O(f)\).+--+-- @since 1.1.0.0+{-# INLINE prod #-}+prod ::+  (HasCallStack, Monoid mono, Monad m) =>+  -- | The `WeightPolicy`.+  WeightPolicy ->+  -- | The `Hld`.+  Hld ->+  -- | User function for getting products in \([u, v)\), where \(u < v\) and+  -- \(\mathrm{depth}(u) < \mathrm{depth}(v)\).+  (VertexHld -> VertexHld -> m mono) ->+  -- | User function for getting products in \([u, v)\), where \(u < v\) and+  -- \(\mathrm{depth}(u) > \mathrm{depth}(v)\).+  (VertexHld -> VertexHld -> m mono) ->+  -- | \(u\).+  Vertex ->+  -- | \(v\).+  Vertex ->+  -- | Product of the path between \(u\) and \(v\).+  m mono+prod weightPolicy hld prodF prodB u0 v0 = do+  foldM+    ( \ !acc (!u, !v) -> do+        !x <-+          if u <= v+            then prodF u $ v + 1+            else prodB v $ u + 1+        pure $! acc <> x+    )+    mempty+    (pathSegmentsInclusive weightPolicy hld u0 v0)

@@ -0,0 +1,297 @@+{-# LANGUAGE RecordWildCards #-}++-- | Integration of segment trees with the heavy-light decomposition technique. Computes monoid+-- products on a path in \(O(\log^2 n)\) time or on a subtree in \(O(\log n)\) time.+--+-- - If vertices have weights, create a `TreeMonoid` with `fromVerts`.+-- - If edges have weights, create a tree monoid with `fromEdges`.+--+-- ==== __(Internals) Weights on edges__+--+-- When vertices are unweighted and only edges have weights, treat edges as new vertices or assign+-- edge weights to the deeper vertex.+--+-- Idea 1. Convert edges into new vertices. This is inefficient.+--+-- @+-- o--o--o  --> o-x-o-x-o+-- @+--+-- Idea 2. Assign edge weight to the deeper vertex. The is the internal implementation of+-- `fromEdges` and LCAs are ignored on `prod`:+--+-- @+--   o+--   | <--- edge 1+--   o <- write weight 1 here+--   | <--- edge 2+--   o <- write weight 2 here+-- @+--+-- ==== __Example (1): Weights are on vertices__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld+-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 0--1--2--3+-- >>> --    ++-- >>> --    +--4--5+-- >>> let n = 6+-- >>> let tree = Gr.build' n . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (2, 3), (1, 4), (4, 5)]+-- >>> let weights = VU.generate n Sum -- vertex `i` is given weight of `i`+-- >>> let hld = Hld.new tree+-- >>> tm <- TM.fromVerts hld {- `Sum` is commutative -} Commute weights+-- >>> TM.prod tm 1 3+-- Sum {getSum = 6}+--+-- >>> TM.prodSubtree tm 1+-- Sum {getSum = 15}+--+-- >>> TM.write tm 1 $ Sum 10+-- >>> TM.prod tm 1 3+-- Sum {getSum = 15}+--+-- ==== __Example (2): Weights are on edges__+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import AtCoder.Extra.Tree.Hld qualified as Hld+-- >>> import AtCoder.Extra.Tree.TreeMonoid qualified as TM+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 0--1--2--3+-- >>> --    ++-- >>> --    +--4--5+-- >>> let n = 6+-- >>> let edges = VU.fromList [(0, 1, Sum (1 :: Int)), (1, 2, Sum 2), (2, 3, Sum 3), (1, 4, Sum 4), (4, 5, Sum 5)]+-- >>> let tree = Gr.build n $ Gr.swapDupe edges+-- >>> let hld = Hld.new tree+-- >>> tm <- TM.fromEdges hld {- `Sum` is commutative -} Commute edges+-- >>> TM.prod tm 1 3+-- Sum {getSum = 5}+--+-- >>> TM.prodSubtree tm 1+-- Sum {getSum = 14}+--+-- >>> TM.write tm 2 $ Sum 10+-- >>> TM.prod tm 1 3+-- Sum {getSum = 13}+--+-- @since 1.1.0.0+module AtCoder.Extra.Tree.TreeMonoid+  ( -- * TreeMonoid+    TreeMonoid,+    Vertex,+    VertexHld,+    Commutativity (..),++    -- * Constructors+    fromVerts,+    fromEdges,++    -- * Segment tree methods++    -- ** Reading+    prod,+    prodSubtree,+    read,++    -- ** Modifications+    write,+    exchange,+    modify,+    modifyM,+  )+where++import AtCoder.Extra.Tree.Hld qualified as Hld+import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.SegTree qualified as ST+import Control.Monad+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Monoid (Dual (..))+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)+import Prelude hiding (read)++-- | Original graph vertex.+--+-- @since 1.1.0.0+type Vertex = Int++-- | Vertex reindexed by `indexHld`.+--+-- @since 1.1.0.0+type VertexHld = Vertex++-- | A wrapper for `Hld` getting product on paths on a tree using `Hld` and segment tree(s).+--+-- @since 1.1.0.0+data TreeMonoid a s = TreeMonoid+  { -- | Borrowed Hld.+    hldTM :: !Hld.Hld,+    -- | Indicates if it's targetting commutative monoids.+    commuteTM :: !Commutativity,+    -- | Indicates if it's targetting edge weights (If not, it's targetting vertex weights).+    weightPolicyTM :: !Hld.WeightPolicy,+    -- | Segment tree for getting products upwards.+    segFTM :: !(ST.SegTree s a),+    -- | Segment tree for getting products downwards. Only created when the monoid is+    -- `NonCommute`.+    segBTM :: !(ST.SegTree s (Dual a))+  }++-- | Represents whether a monoid is commutative or noncommutative.+--+-- @since 1.1.0.0+data Commutativity+  = -- | Commutative: \(a \cdot b = b \cdot a\).+    --+    -- @since 1.1.0.0+    Commute+  | -- | Noncommutative: \(a \cdot b \neq b \cdot a\).+    --+    -- @since 1.1.0.0+    NonCommute+  deriving+    ( -- | @since 1.1.0.0+      Eq,+      -- | @since 1.1.0.0+      Show+    )++-- | \(O(n)\)+buildImpl ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  Hld.Hld ->+  Commutativity ->+  Hld.WeightPolicy ->+  VU.Vector a ->+  m (TreeMonoid a (PrimState m))+buildImpl hldTM commuteTM weightPolicyTM weights = do+  segFTM <- ST.build weights+  segBTM <-+    case commuteTM of+      Commute -> ST.build VU.empty+      NonCommute -> ST.build $ VU.map Dual weights+  pure TreeMonoid {..}++-- | \(O(n)\) Creates a `TreeMonoid` with weights on vertices.+--+-- @since 1.1.0.0+fromVerts ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | `Hld.Hld`.+  Hld.Hld ->+  -- | Whether the monoid is commutative or not.+  Commutativity ->+  -- | The vertex weights.+  VU.Vector a ->+  -- | A `TreeMonoid` with weights on vertices.+  m (TreeMonoid a (PrimState m))+fromVerts hld@Hld.Hld {indexHld} commuteTM xs_ = do+  let !_ = ACIA.runtimeAssert (VU.length indexHld == VU.length xs_) $ "AtCoder.Extra.Tree.TreeMonoid.fromVerts: vertex number mismatch (`" ++ show (VU.length indexHld) ++ "` and `" ++ show (VU.length xs_) ++ "`)"+  let !xs = VU.create $ do+        vec <- VUM.unsafeNew $ VU.length xs_+        VU.iforM_ xs_ $ \i x -> do+          VGM.write vec (indexHld VG.! i) x+        pure vec+  buildImpl hld commuteTM Hld.WeightsAreOnVertices xs++-- | \(O(n)\) Creates a `TreeMonoid` with weignts on edges. The edges are not required to be+-- duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is needed.+--+-- @since 1.1.0.0+fromEdges ::+  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | `Hld.Hld`.+  Hld.Hld ->+  -- | Whether the monoid is commutative or not.+  Commutativity ->+  -- | Input edges.+  VU.Vector (Vertex, Vertex, a) ->+  -- | A `TreeMonoid` with weights on edges.+  m (TreeMonoid a (PrimState m))+fromEdges hld@Hld.Hld {indexHld} commuteTM edges = do+  let !xs = VU.create $ do+        vec <- VUM.unsafeNew $ VU.length indexHld+        VU.forM_ edges $ \(!u, !v, !w) -> do+          let u' = indexHld VG.! u+          let v' = indexHld VG.! v+          VGM.write vec (max u' v') w+        pure vec+  buildImpl hld commuteTM Hld.WeightsAreOnEdges xs++-- | \(O(\log^2 n)\) Returns the product of the path between two vertices \(u\), \(v\) (invlusive).+--+-- @since 1.1.0.0+prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> Vertex -> m a+prod TreeMonoid {..} u v = do+  case commuteTM of+    Commute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (ST.prod segFTM) u v+    NonCommute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (\l r -> getDual <$> ST.prod segBTM l r) u v++-- | \(O(\log n)\) Returns the product of the subtree rooted at the given `Vertex`.+--+-- @since 1.1.0.0+prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a+prodSubtree TreeMonoid {..} subtreeRoot = do+  let (!l, !r) = Hld.subtreeSegmentInclusive hldTM subtreeRoot+  case weightPolicyTM of+    Hld.WeightsAreOnVertices -> ST.prod segFTM l (r + 1)+    Hld.WeightsAreOnEdges -> do+      -- ignore the root of the subtree, which has the minimum index among the subtree vertices+      if l == r+        then pure mempty+        else ST.prod segFTM (l + 1) (r + 1)++-- | \(O(1)\) Reads a `TreeMonoid` value on a `Vertex`.+--+-- @since 1.1.0.0+read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a+read TreeMonoid {..} i_ = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.read segFTM i++-- | \(O(\log n)\) Write a `TreeMonoid` value on a `Vertex`.+--+-- @since 1.1.0.0+write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m ()+write TreeMonoid {..} i_ x = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.write segFTM i x+  when (commuteTM == NonCommute) $ do+    ST.write segBTM i $ Dual x++-- | \(O(\log n)\) Exchanges a `TreeMonoid` value on a `Vertex`.+--+-- @since 1.1.0.0+exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m a+exchange TreeMonoid {..} i_ x = do+  let !i = Hld.indexHld hldTM VG.! i_+  !res <- ST.exchange segFTM i x+  when (commuteTM == NonCommute) $ do+    ST.write segBTM i $ Dual x+  pure res++-- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`.+--+-- @since 1.1.0.0+modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()+modify TreeMonoid {..} f i_ = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.modify segFTM f i+  when (commuteTM == NonCommute) $ do+    ST.modify segBTM (Dual . f . getDual) i++-- | \(O(\log n)\) Modifies a `TreeMonoid` value on a `Vertex`.+--+-- @since 1.1.0.0+modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m ()+modifyM TreeMonoid {..} f i_ = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.modifyM segFTM f i+  when (commuteTM == NonCommute) $ do+    ST.modifyM segBTM ((Dual <$>) . f . getDual) i

@@ -0,0 +1,407 @@+{-# LANGUAGE RecordWildCards #-}++-- | A static Wavelet Matrix.+--+-- ==== Notation+-- Let \(S\) be the set of values in your wavelet matrix. We use \(|S|\) to denote the number of+-- distinct values contained within this set \((|S| \lt n)\).+--+-- @since 1.1.0.0+module AtCoder.Extra.WaveletMatrix+  ( -- * Wavelet Matrix+    WaveletMatrix (..),++    -- * Constructors+    build,++    -- * Access (indexing)+    access,++    -- * Rank (count)+    rank,+    rankBetween,++    -- * Selection++    -- | ==== __Example__+    -- >>> import AtCoder.Extra.WaveletMatrix qualified as WM+    -- >>> import Data.Vector.Unboxed qualified as VU+    -- >>> let wm = WM.build $ VU.fromList [1,1,2,1,3]+    -- >>> WM.select wm 1+    -- Just 0+    -- >>> WM.selectKth wm 2 1+    -- Just 3+    -- >>> WM.selectIn wm {- [l, r) -} 1 4 {- x -} 1+    -- Just 1+    -- >>> WM.selectKthIn wm {- [l, r) -} 1 4 {- k -} 1 {- x -} 1+    -- Just 3+    select,+    selectKth,+    selectIn,+    selectKthIn,++    -- * Quantile (value-ordered access)+    kthLargestIn,+    ikthLargestIn,+    kthSmallestIn,+    ikthSmallestIn,+    -- unsafeKthLargestIn,+    -- unsafeIKthLargestIn,+    -- unsafeKthSmallestIn,+    -- unsafeIKthSmallestIn,++    -- * Lookup+    lookupLE,+    lookupLT,+    lookupGE,+    lookupGT,++    -- * Conversions+    assocsIn,+    descAssocsIn,+  )+where++import AtCoder.Extra.Bisect+import AtCoder.Extra.WaveletMatrix.Raw qualified as Rwm+import Control.Monad+import Data.Maybe (fromJust, fromMaybe)+import Data.Vector.Algorithms.Intro qualified as VAI+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU++-- | A static Wavelet Matrix.+--+-- @since 1.1.0.0+data WaveletMatrix = WaveletMatrix+  { -- | The internal wavelet matrix, where index compression is not automatically performed.+    --+    -- @since 1.1.0.0+    rawWM :: !Rwm.RawWaveletMatrix,+    -- | Index compression dictionary.+    --+    -- @since 1.1.0.0+    xDictWM :: !(VU.Vector Int)+  }++-- | \(O(n \log n)\) Creates a `WaveletMatrix` from an array \(a\).+--+-- @since 1.1.0.0+{-# INLINE build #-}+build :: VU.Vector Int -> WaveletMatrix+build ys =+  let !xDictWM = VU.uniq $ VU.modify (VAI.sortBy compare) ys+      !ys' = VU.map (fromJust . lowerBound xDictWM) ys+      !rawWM = Rwm.build (VG.length ys) ys'+   in WaveletMatrix {..}++-- | \(O(\log |S|)\) Returns \(a[k]\) or `Nothing` if the index is out of the bounds. Try to use the+-- original array if you can.+--+-- @since 1.1.0.0+{-# INLINE access #-}+access :: WaveletMatrix -> Int -> Maybe Int+access WaveletMatrix {..} i = (xDictWM VG.!) <$> Rwm.access rawWM i++-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\).+--+-- @since 1.1.0.0+{-# INLINE rank #-}+rank ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The number of \(y\) in \([l, r)\).+  Int+rank wm l r y = rankBetween wm l r y (y + 1)++-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\).+--+-- @since 1.1.0.0+{-# INLINE rankBetween #-}+rankBetween ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_1\)+  Int ->+  -- | \(y_2\)+  Int ->+  -- | The number of \(y\) in \([l, r) \times [y_1, y_2)\).+  Int+rankBetween WaveletMatrix {..} l r y1 y2+  | not $ 0 <= l && l < r && r <= n = 0+  | y1' >= y2' = 0+  | otherwise = Rwm.rankBetween rawWM l r y1' y2'+  where+    -- Handles the case @yl@ or  @yr@ is not in the dict+    n = Rwm.lengthRwm rawWM+    y1' = fromMaybe n (bisectR 0 (VG.length xDictWM) ((< y1) . VG.unsafeIndex xDictWM))+    y2' = maybe (-1) (+ 1) (bisectL 0 (VG.length xDictWM) ((< y2) . VG.unsafeIndex xDictWM))++-- | \(O(\log |S|)\) Returns the index of the first \(y\) in \(a\), or `Nothing` if \(y\) is+-- not found.+--+-- @since 1.1.0.0+{-# INLINE select #-}+select :: WaveletMatrix -> Int -> Maybe Int+select wm = selectKth wm 0++-- | \(O(\log |S|)\) Returns the index of the \(k\)-th occurrence (0-based) of \(y\), or `Nothing`+-- if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectKth #-}+selectKth ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(k\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The index of \(k\)-th \(y\)+  Maybe Int+selectKth WaveletMatrix {..} k y = do+  i <- lowerBound xDictWM y+  -- TODO: we don't need such an explicit branch?+  let !y' = xDictWM VG.! i+  guard $ y' == y+  Rwm.selectKth rawWM k i++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the first occurrence+-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectIn #-}+selectIn ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The index of the first \(y\) in \([l, r)\).+  Maybe Int+selectIn wm l r = selectKthIn wm l r 0++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th occurrence+-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectKthIn #-}+selectKthIn ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The index of the \(k\)-th \(y\) in \([l, r)\).+  Maybe Int+selectKthIn WaveletMatrix {..} l r k y = do+  i <- lowerBound xDictWM y+  -- TODO: we don't need such an explicit branch?+  let !y' = xDictWM VG.! i+  guard $ y' == y+  Rwm.selectKthIn rawWM l r k i++-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based)+-- largest value. Note that duplicated values are treated as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE kthLargestIn #-}+kthLargestIn ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(k\)-th largest \(y\) in \([l, r)\)+  Maybe Int+kthLargestIn WaveletMatrix {..} l r k+  | Just !y <- Rwm.kthLargestIn rawWM l r k = Just $ xDictWM VG.! y+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the+-- \(k\)-th (0-based) largest value. Note that duplicated values are treated as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE ikthLargestIn #-}+ikthLargestIn ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)+  Maybe (Int, Int)+ikthLargestIn WaveletMatrix {..} l r k+  | Just (!i, !y) <- Rwm.ikthLargestIn rawWM l r k = Just (i, xDictWM VG.! y)+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns the index of the \(k\)-th (0-based)+-- smallest value. Note that duplicated values are treated as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE kthSmallestIn #-}+kthSmallestIn ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(k\)-th largest \(y\) in \([l, r)\)+  Maybe Int+kthSmallestIn WaveletMatrix {..} l r k+  | Just !y <- Rwm.kthSmallestIn rawWM l r k = Just $ xDictWM VG.! y+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given the interval \([l, r)\), returns both the index and the value of the+-- \(k\)-th (0-based) smallest value. Note that duplicated values are treated as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE ikthSmallestIn #-}+ikthSmallestIn ::+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)+  Maybe (Int, Int)+ikthSmallestIn WaveletMatrix {..} l r k+  | Just (!i, !y) <- Rwm.ikthSmallestIn rawWM l r k = Just (i, xDictWM VG.! y)+  | otherwise = Nothing++-- | \(O(\log |S|)\)+--+-- @since 1.1.0.0+{-# INLINE unsafeKthSmallestIn #-}+unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int+unsafeKthSmallestIn WaveletMatrix {..} l r k =+  xDictWM VG.! Rwm.unsafeKthSmallestIn rawWM l r k++-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\).+--+-- @since 1.1.0.0+{-# INLINE lookupLE #-}+lookupLE ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0]\)+  Maybe Int+lookupLE wm l r y0+  | r' == l' = Nothing+  | rank_ == 0 = Nothing+  | otherwise = Just $ unsafeKthSmallestIn wm l' r' (rank_ - 1)+  where+    -- clamp+    l' = max 0 l+    r' = min (Rwm.lengthRwm (rawWM wm)) r+    rank_ = rankBetween wm l' r' minBound (y0 + 1)++-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\).+--+-- @since 1.1.0.0+{-# INLINE lookupLT #-}+lookupLT ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0)\)+  Maybe Int+lookupLT wm l r y0 = lookupLE wm l r (y0 - 1)++-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\).+--+-- @since 1.1.0.0+{-# INLINE lookupGE #-}+lookupGE ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Minimum \(y\) in \([l, r) \times [y_0, \infty)\).+  Maybe Int+lookupGE wm l r y0+  | r' == l' = Nothing+  | rank_ >= r - l = Nothing+  | otherwise = Just $ unsafeKthSmallestIn wm l r rank_+  where+    -- clamp+    l' = max 0 l+    r' = min (Rwm.lengthRwm (rawWM wm)) r+    rank_ = rankBetween wm l' r' minBound y0++-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\).+--+-- @since 1.1.0.0+{-# INLINE lookupGT #-}+lookupGT ::+  -- | A wavelet matrix+  WaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Minimum \(y\) in \([l, r) \times (y_0, \infty)\)+  Maybe Int+lookupGT wm l r y0 = lookupGE wm l r (y0 + 1)++-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in+-- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small.+--+-- @since 1.1.0.0+{-# INLINE assocsIn #-}+assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]+assocsIn WaveletMatrix {..} l r = Rwm.assocsWith rawWM l r (xDictWM VG.!)++-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in+-- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small.+--+-- @since 1.1.0.0+{-# INLINE descAssocsIn #-}+descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)]+descAssocsIn WaveletMatrix {..} l r = Rwm.descAssocsInWith rawWM l r (xDictWM VG.!)

@@ -0,0 +1,179 @@+{-# LANGUAGE RecordWildCards #-}++-- | A compact bit vector, a collection of bits that can process @rank@ in \(O(1)\) and @select@ in+-- \(O(\log n)\).+--+-- @since 1.1.0.0+module AtCoder.Extra.WaveletMatrix.BitVector+  ( -- * Bit vector+    BitVector (..),++    -- * Constructor+    build,++    -- * (Internal) Word-based cumultaive sum+    wordSize,+    csumInPlace,++    -- * Rank+    rank0,+    rank1,++    -- * Select+    select0,+    select1,+    selectKthIn0,+    selectKthIn1,+  )+where++import AtCoder.Extra.Bisect (bisectL)+import Control.Monad.Primitive (PrimMonad (PrimState))+import Data.Bit (Bit (..))+import Data.Bits (popCount)+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM++-- | A compact bit vector.+--+-- @since 1.1.0.0+data BitVector = BitVector+  { -- | Packed bits.+    --+    -- @since 1.1.0.0+    bitsBv :: !(VU.Vector Bit),+    -- | Cumulative sum of bits by 64 words.+    --+    -- @since 1.1.0.0+    csumBv :: !(VU.Vector Int)+    -- we could use Word32 for csumBv, as 2^32 is large enough+  }+  deriving (Eq, Show)++-- | \(O(n)\) Creates a `BitVector`.+--+-- @since 1.1.0.0+{-# INLINE build #-}+build :: VU.Vector Bit -> BitVector+build bitsBv =+  let csumBv = VU.create $ do+        vec <- VUM.replicate ((VU.length bitsBv + 63) `div` 64 + 1) 0+        _ <- csumInPlace vec bitsBv+        pure vec+   in BitVector {..}++-- | The block size \(64\) for the internal cumultaive sum in the bit vector.+--+-- @since 1.1.0.0+{-# INLINE wordSize #-}+wordSize :: Int+wordSize = 64++-- | \(O(n)\) Calculates the cumulative sum in-place for the bit vector and returns the sum.+--+-- @since 1.1.0.0+{-# INLINE csumInPlace #-}+csumInPlace ::+  (PrimMonad m) =>+  -- | Cumulative sum of length \(\lceil |\mathrm{bits}| / 64 \rceil\).+  VUM.MVector (PrimState m) Int ->+  -- | Bits+  VU.Vector Bit ->+  -- | Sum of the bits+  m Int+csumInPlace csum bits = do+  VGM.unsafeWrite csum 0 (0 :: Int)++  -- Calcuate popCount by word. TODO: use `castToWords` for most elements+  VU.ifoldM'+    ( \ !acc i wordSum -> do+        let !acc' = acc + wordSum+        VGM.unsafeWrite csum (i + 1) acc'+        pure acc'+    )+    (0 :: Int)+    $ VU.unfoldrExactN+      (VGM.length csum - 1)+      (\bits' -> (popCount (VU.take wordSize bits'), VU.drop wordSize bits'))+      bits++-- | \(O(1)\) Counts the number of \(0\) bits in the interval \([0, i)\).+--+-- @since 1.1.0.0+{-# INLINE rank0 #-}+rank0 :: BitVector -> Int -> Int+rank0 bv i = i - rank1 bv i++-- | \(O(1)\) Counts the number of \(1\) bits in an interval \([0, i)\).+--+-- @since 1.1.0.0+{-# INLINE rank1 #-}+rank1 :: BitVector -> Int -> Int+rank1 BitVector {..} i = fromCSum + fromRest+  where+    -- TODO: check bounds for i?+    (!nWords, !nRest) = i `divMod` wordSize+    fromCSum = VG.unsafeIndex csumBv nWords+    fromRest = popCount . VU.take nRest . VU.drop (nWords * wordSize) $ bitsBv++-- | \(O(\log n)\) Returns the index of \(k\)-th \(0\) (0-based), or `Nothing` if no such bit exists.+--+-- @since 1.1.0.0+{-# INLINE select0 #-}+select0 :: BitVector -> Int -> Maybe Int+select0 bv = selectKthIn0 bv 0 (VG.length (bitsBv bv))++-- | \(O(\log n)\) Returns the index of \(k\)-th \(1\) (0-based), or `Nothing` if no such bit exists.+--+-- @since 1.1.0.0+{-# INLINE select1 #-}+select1 :: BitVector -> Int -> Maybe Int+select1 bv = selectKthIn1 bv 0 (VG.length (bitsBv bv))++-- | \(O(\log n)\) Returns the index of \(k\)-th \(0\) (0-based) in \([l, r)\), or `Nothing` if no+-- such bit exists.+--+-- @since 1.1.0.0+{-# INLINE selectKthIn0 #-}+selectKthIn0 ::+  -- | A bit vector+  BitVector ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | The index of \(k\)-th \(0\) in \([l, r)\)+  Maybe Int+selectKthIn0 bv l r k+  | k < 0 || nZeros <= k = Nothing+  | otherwise = bisectL l r $ \i -> rank0 bv i - rankL0 < k + 1+  where+    nZeros = rank0 bv r - rankL0+    rankL0 = rank0 bv l++-- | \(O(\log n)\) Returns the index of \(k\)-th \(1\) (0-based) in \([l, r)\), or `Nothing` if no+-- such bit exists.+--+-- @since 1.1.0.0+{-# INLINE selectKthIn1 #-}+selectKthIn1 ::+  -- | A bit vector+  BitVector ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | The index of \(k\)-th \(1\) in \([l, r)\)+  Maybe Int+selectKthIn1 bv l r k+  | k < 0 || nOnes <= k = Nothing+  | otherwise = bisectL l r $ \i -> rank1 bv i - rankL1 < k + 1+  where+    nOnes = rank1 bv r - rankL1+    rankL1 = rank1 bv l

@@ -0,0 +1,651 @@+{-# LANGUAGE RecordWildCards #-}++-- original implementation:+-- <https://miti-7.hatenablog.com/entry/2018/04/28/152259>++-- NOTE: We could integrate cumulative sum / fenwick tree / segment tree.+-- NOTE: `topK` and `intersects` are not implemented as they are slow.++-- | A static Wavelet Matrix without automatic index comperssion. Consider using+-- @AtCoder.Extra.WaveletMatrix@ instead.+--+-- @since 1.1.0.0+module AtCoder.Extra.WaveletMatrix.Raw+  ( -- * RawWaveletMatrix+    RawWaveletMatrix (..),++    -- * Constructors+    build,++    -- * Access (indexing)+    access,++    -- * rank+    rankLT,+    rank,+    rankBetween,++    -- * Select+    select,+    selectKth,+    selectIn,+    selectKthIn,++    -- * Quantile (value-ordered access)++    -- ** Safe (total)+    kthSmallestIn,+    ikthSmallestIn,+    kthLargestIn,+    ikthLargestIn,++    -- ** Unsafe+    unsafeKthSmallestIn,+    unsafeIKthSmallestIn,+    unsafeKthLargestIn,+    unsafeIKthLargestIn,++    -- * Lookup+    lookupLE,+    lookupLT,+    lookupGE,+    lookupGT,++    -- * Conversions+    assocsIn,+    assocsWith,+    descAssocsIn,+    descAssocsInWith,+  )+where++import AtCoder.Extra.WaveletMatrix.BitVector qualified as BV+import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.Internal.Bit qualified as ACIB+import Control.Monad.ST (runST)+import Data.Bit (Bit (..))+import Data.Bits (bit, countTrailingZeros, setBit, testBit, (.|.))+import Data.Maybe+import Data.Vector qualified as V+import Data.Vector.Algorithms.Radix qualified as VAR+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++-- | A static Wavelet Matrix without automatic index comperssion.+--+-- @since 1.1.0.0+data RawWaveletMatrix = RawWaveletMatrix+  { -- | \(\lceil \log_2 N \rceil\).+    --+    -- @since 1.1.0.0+    heightRwm :: {-# UNPACK #-} !Int,+    -- | The length of the original array.+    --+    -- @since 1.1.0.0+    lengthRwm :: {-# UNPACK #-} !Int,+    -- | The bit matrix. Each row represents (heightRwm - 1 - iRow) bit's on/off.+    --+    -- @since 1.1.0.0+    bitsRwm :: !(V.Vector BV.BitVector),+    -- | The number of zeros bits in each row in the bit matrix.+    --+    -- @since 1.1.0.0+    nZerosRwm :: !(VU.Vector Int)+  }+  deriving (Eq, Show)++-- | \(O(n \log n)\) Creates a `RawWaveletMatrix` from a vector \(a\).+--+-- @since 1.1.0.0+{-# INLINE build #-}+build ::+  (HasCallStack) =>+  -- | The number of different values in the compressed vector.+  Int ->+  -- | A compressed vector+  VU.Vector Int ->+  -- | A wavelet matrix+  RawWaveletMatrix+build nx xs+  | nx < 0 = error "AtCoder.Extra.WaveletMatrix.Raw.build: given negative `n`"+  | otherwise = runST $ do+      -- TODO: less mutable variables+      orgBits <- VUM.replicate (lengthRwm * heightRwm) $ Bit False+      orgCsum <- VUM.replicate (lenCSum * heightRwm) (0 :: Int)+      nZeros <- VUM.unsafeNew heightRwm++      -- views by row over the contiguous memory:+      let !bits = V.unfoldrExactN heightRwm (VUM.splitAt lengthRwm) orgBits+      let !csums = V.unfoldrExactN heightRwm (VUM.splitAt lenCSum) orgCsum++      -- the vector will be sorted by bits.+      vec <- VU.thaw xs+      V.izipWithM_+        ( \iRow bitVec csum -> do+            let !iBit = heightRwm - 1 - iRow+            vec' <- VU.unsafeFreeze vec+            VU.iforM_ vec' $ \i x -> do+              VGM.unsafeWrite bitVec i . Bit $ testBit x iBit++            -- csum.+            VGM.unsafeWrite csum 0 (0 :: Int)+            bitVec' <- VU.unsafeFreeze bitVec++            -- get popCount by word. TODO: use `castToWords` for most elements+            nOnes <- BV.csumInPlace csum bitVec'+            VGM.unsafeWrite nZeros iRow (lengthRwm - nOnes)++            -- preform a stable sort by the bit:+            VAR.sortBy 2 2 (\_ x -> fromEnum (testBit x iBit)) vec+        )+        bits+        csums+      nZerosRwm <- VU.unsafeFreeze nZeros+      bits' <- V.unfoldrExactN heightRwm (VU.splitAt lengthRwm) <$> VU.unsafeFreeze orgBits+      csums' <- V.unfoldrExactN heightRwm (VU.splitAt lenCSum) <$> VU.unsafeFreeze orgCsum+      let !bitsRwm = V.zipWith BV.BitVector bits' csums'+      pure $ RawWaveletMatrix {..}+  where+    !lengthRwm = VG.length xs+    !lenCSum = (lengthRwm + BV.wordSize - 1) `div` BV.wordSize + 1 -- +1 for the zero+    !heightRwm = countTrailingZeros $ ACIB.bitCeil nx++-- | \(O(\log |S|)\) Returns \(a[k]\) or `Nothing` if the index is out of the bounds. Try to use the+-- original array if you can.+--+-- @since 1.1.0.0+{-# INLINE access #-}+access :: RawWaveletMatrix -> Int -> Maybe Int+access RawWaveletMatrix {..} i0+  | ACIA.testIndex i0 lengthRwm =+      let (!_, !res) =+            V.ifoldl'+              ( \(!i, !acc) !iRow !bits ->+                  let Bit !goRight = VG.unsafeIndex (BV.bitsBv bits) i+                      !i'+                        | goRight = BV.rank1 bits i + VG.unsafeIndex nZerosRwm iRow+                        | otherwise = BV.rank0 bits i+                      !acc'+                        | goRight = setBit acc (heightRwm - 1 - iRow)+                        | otherwise = acc+                   in (i', acc')+              )+              (i0, 0)+              bitsRwm+       in Just res+  | otherwise = Nothing++-- | \(O(\log |A|)\) Goes down the wavelet matrix for collecting the kth smallest value.+--+-- @since 1.1.0.0+{-# INLINE goDown #-}+goDown :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int, Int, Int)+goDown RawWaveletMatrix {..} l_ r_ k_ = V.ifoldl' step (0 :: Int, l_, r_, k_) bitsRwm+  where+    -- It's binary search over the value range. In each row, we'll focus on either 0 bit values or+    -- 1 bit values in [l, r) and update the range to [l', r').+    step (!acc, !l, !r, !k) !iRow !bits+      -- `r0 - l0`, the number of zeros in [l, r), is bigger than or equal to k:+      -- Go left.+      | k < r0 - l0 = (acc, l0, r0, k)+      -- Go right.+      | otherwise =+          let !acc' = acc .|. bit (heightRwm - 1 - iRow)+              !nZeros = VG.unsafeIndex nZerosRwm iRow+              -- every zero bits come to the left after the move.+              !l' = l + nZeros - l0 -- add the number of zeros in [0, l)+              !r' = r + nZeros - r0 -- add the number of zeros in [0, r)+              !k' = k - (r0 - l0) -- `r0 - l0` zeros go left+           in (acc', l', r', k')+      where+        !l0 = BV.rank0 bits l+        !r0 = BV.rank0 bits r++-- | \(O(\log |A|)\) Goes up the wavelet matrix for collecting the value \(x\).+--+-- @since 1.1.0.0+{-# INLINE goUp #-}+goUp :: RawWaveletMatrix -> Int -> Int -> Maybe Int+goUp RawWaveletMatrix {..} i0 x =+  V.ifoldM'+    ( \ !i !iBit !bits ->+        if testBit x iBit+          then BV.select1 bits $ i - nZerosRwm VG.! (heightRwm - 1 - iBit)+          else BV.select0 bits i+    )+    i0+    (V.reverse bitsRwm)++-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [0, y_0)\).+--+-- @since 1.1.0.0+{-# INLINE rankLT #-}+rankLT :: RawWaveletMatrix -> Int -> Int -> Int -> Int+rankLT RawWaveletMatrix {..} l_ r_ xr+  -- REMARK: This is required. The function below cannot handle the case N = 2^i and xr = N.+  | xr >= bit heightRwm = r'_ - l'_+  | xr <= 0 = 0+  | r'_ <= l'_ = 0+  | otherwise =+      let (!res, !_, !_) = V.ifoldl' step (0, l'_, r'_) bitsRwm+       in res+  where+    -- clamp+    l'_ = max 0 l_+    r'_ = min lengthRwm r_+    -- [l, r)+    step (!acc, !l, !r) !iRow !bits =+      let !b = testBit xr (heightRwm - 1 - iRow)+          !l0 = BV.rank0 bits l+          !r0 = BV.rank0 bits r+       in if b+            then (acc + r0 - l0, l - l0 + VG.unsafeIndex nZerosRwm iRow, r - r0 + VG.unsafeIndex nZerosRwm iRow)+            else (acc, l0, r0)++-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\).+--+-- @since 1.1.0.0+{-# INLINE rank #-}+rank ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The number of \(y\) in \([l, r)\).+  Int+rank wm l r x = rankBetween wm l r x (x + 1)++-- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\).+--+-- @since 1.1.0.0+{-# INLINE rankBetween #-}+rankBetween ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_1\)+  Int ->+  -- | \(y_2\)+  Int ->+  -- | The number of \(y\) in \([l, r) \times [y_1, y_2)\).+  Int+rankBetween wm l r lx rx = rankLT wm l r rx - rankLT wm l r lx++-- | \(O(\log |S|)\) Returns the index of the first \(y\) in the sequence, or `Nothing` if \(y\) is+-- not found.+--+-- @since 1.1.0.0+{-# INLINE select #-}+select :: RawWaveletMatrix -> Int -> Maybe Int+select wm = selectKth wm 0++-- | \(O(\log |S|)\) Returns the index of the \(k\)-th occurrence (0-based) of \(y\), or `Nothing`+-- if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectKth #-}+selectKth ::+  RawWaveletMatrix ->+  -- | \(k\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The index of \(k\)-th \(y\)+  Maybe Int+selectKth wm = selectKthIn wm 0 (lengthRwm wm)++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the first occurrence+-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectIn #-}+selectIn ::+  -- | A wavelet matrix+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | The index of the first \(y\) in \([l, r)\).+  Maybe Int+selectIn wm = selectKthIn wm 0++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th occurrence+-- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists.+--+-- @since 1.1.0.0+{-# INLINE selectKthIn #-}+selectKthIn ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(y\)+  Int ->+  -- | The index of the \(k\)-th \(y\) in \([l, r)\).+  Maybe Int+selectKthIn wm@RawWaveletMatrix {..} l_ r_ k x+  | not (0 <= x && x < lengthRwm && 0 <= k && k < lengthRwm) = Nothing+  | l'_ < r'_ = inner+  | otherwise = Nothing+  where+    -- clamp+    l'_ = max 0 l_+    r'_ = min lengthRwm r_+    inner :: Maybe Int+    inner+      | rEnd <= lEnd + k = Nothing+      -- go up+      | otherwise = goUp wm (lEnd + k) x+      where+        -- TODO: replace with goDown+        -- Go down. Gets the [l, r) range of @x@ in the last array.+        (!lEnd, !rEnd) =+          V.ifoldl'+            ( \(!l, !r) !iRow !bits ->+                let !l0 = BV.rank0 bits l+                    !r0 = BV.rank0 bits r+                 in if testBit x (heightRwm - 1 - iRow)+                      then (l + nZerosRwm VG.! iRow - l0, r + nZerosRwm VG.! iRow - r0)+                      else (l0, r0)+            )+            (l'_, r'_)+            bitsRwm++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th (0-based)+-- largest value. Note that duplicated values are counted as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE kthLargestIn #-}+kthLargestIn ::+  -- | A wavelet matrix+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(k\)-th largest \(y\) in \([l, r)\)+  Maybe Int+kthLargestIn wm l r k+  | k < 0 || k >= r - l = Nothing+  | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeKthLargestIn wm l r k+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns both the index and the value of the+-- \(k\)-th (0-based) largest value. Note that duplicated values are counted as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE ikthLargestIn #-}+ikthLargestIn ::+  -- | A wavelet matrix+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)+  Maybe (Int, Int)+ikthLargestIn wm l r k+  | k < 0 || k >= r - l = Nothing+  | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeIKthLargestIn wm l r k+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns the index of the \(k\)-th (0-based)+-- smallest value. Note that duplicated values are counted as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE kthSmallestIn #-}+kthSmallestIn ::+  -- | A wavelet matrix+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \(k\)-th largest \(y\) in \([l, r)\)+  Maybe Int+kthSmallestIn wm l r k+  | k < 0 || k >= r - l = Nothing+  | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeKthSmallestIn wm l r k+  | otherwise = Nothing++-- | \(O(\log |S|)\) Given an interval \([l, r)\), it returns both the index and the value of the+-- \(k\)-th (0-based) smallest value. Note that duplicated values are counted as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE ikthSmallestIn #-}+ikthSmallestIn ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(k\)+  Int ->+  -- | \((i, y)\) for \(k\)-th largest \(y\) in \([l, r)\)+  Maybe (Int, Int)+ikthSmallestIn wm l r k+  | k < 0 || k >= r - l = Nothing+  | 0 <= l && l < r && r <= lengthRwm wm = Just $ unsafeIKthSmallestIn wm l r k+  | otherwise = Nothing++-- | \(O(\log a)\) Returns \(k\)-th (0-based) biggest number in \([l, r)\). Note that duplicated+-- values are counted as distinct occurrences.+--+-- @since 1.1.0.0+{-# INLINE unsafeKthLargestIn #-}+unsafeKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int+unsafeKthLargestIn wm l r k = unsafeKthSmallestIn wm l r (r - l - (k + 1))++-- | \(O(\log a)\)+--+-- @since 1.1.0.0+{-# INLINE unsafeIKthLargestIn #-}+unsafeIKthLargestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int)+unsafeIKthLargestIn wm l r k = unsafeIKthSmallestIn wm l r (r - l - (k + 1))++-- | \(O(\log a)\)+--+-- @since 1.1.0.0+{-# INLINE unsafeKthSmallestIn #-}+unsafeKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> Int+unsafeKthSmallestIn wm l_ r_ k_ =+  let (!x, !_, !_, !_) = goDown wm l_ r_ k_+   in x++-- | \(O(\log a)\)+--+-- @since 1.1.0.0+{-# INLINE unsafeIKthSmallestIn #-}+unsafeIKthSmallestIn :: RawWaveletMatrix -> Int -> Int -> Int -> (Int, Int)+unsafeIKthSmallestIn wm l_ r_ k_ =+  let (!x, !l, !_, !k) = goDown wm l_ r_ k_+      !i' = fromJust $ goUp wm (l + k) x+   in (i', x)++-- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\).+--+-- @since 1.1.0.0+{-# INLINE lookupLE #-}+lookupLE ::+  -- | A wavelet matrix+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0]\)+  Maybe Int+lookupLE wm l r x+  | r' == l' = Nothing+  | rank_ == 0 = Nothing+  | otherwise = Just $ unsafeKthSmallestIn wm l' r' (rank_ - 1)+  where+    -- clamp+    l' = max 0 l+    r' = min (lengthRwm wm) r+    rank_ = rankBetween wm l r minBound (x + 1)++-- | \(O(\log a)\) Finds the maximum \(x\) in \([l, r)\) s.t. \(x_{0} \lt x\).+--+-- @since 1.1.0.0+{-# INLINE lookupLT #-}+lookupLT ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(x\)+  Int ->+  -- | Maximum \(y\) in \([l, r) \times (-\infty, y_0)\)+  Maybe Int+lookupLT wm l r x = lookupLE wm l r (x - 1)++-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\).+--+-- @since 1.1.0.0+{-# INLINE lookupGE #-}+lookupGE ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Minimum \(y\) in \([l, r) \times [y_0, \infty)\).+  Maybe Int+lookupGE wm l r x+  | r' == l' = Nothing+  | rank_ >= r' - l' = Nothing+  | otherwise =+      Just $ unsafeKthSmallestIn wm l' r' rank_+  where+    -- clamp+    l' = max 0 l+    r' = min (lengthRwm wm) r+    rank_ = rankBetween wm l' r' minBound x++-- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\).+--+-- @since 1.1.0.0+{-# INLINE lookupGT #-}+lookupGT ::+  RawWaveletMatrix ->+  -- | \(l\)+  Int ->+  -- | \(r\)+  Int ->+  -- | \(y_0\)+  Int ->+  -- | Minimum \(y\) in \([l, r) \times (y_0, \infty)\)+  Maybe Int+lookupGT wm l r x = lookupGE wm l r (x + 1)++-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in+-- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small.+--+-- @since 1.1.0.0+{-# INLINE assocsIn #-}+assocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)]+assocsIn wm l r = assocsWith wm l r id++-- | \(O(\log A \min(|A|, L))\) Internal implementation of `assocs`.+--+-- @since 1.1.0.0+{-# INLINE assocsWith #-}+assocsWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)]+assocsWith RawWaveletMatrix {..} l_ r_ f+  | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ []+  | otherwise = []+  where+    -- clamp+    l'_ = max 0 l_+    r'_ = min lengthRwm r_+    -- DFS. [l, r)+    inner !acc iRow !l !r res+      | iRow >= heightRwm =+          let !n = r - l+              !acc' = f acc+           in (acc', n) : res+      | otherwise = do+          let !bits = bitsRwm VG.! iRow+              !l0 = BV.rank0 bits l+              !r0 = BV.rank0 bits r+              !nZeros = nZerosRwm VG.! iRow+              -- go right (visit bigger values first)+              !l' = l + nZeros - l0+              !r' = r + nZeros - r0+              !res'+                | l' < r' = inner (acc .|. bit (heightRwm - 1 - iRow)) (iRow + 1) l' r' res+                | otherwise = res+              !res''+                -- go left+                | l0 < r0 = inner acc (iRow + 1) l0 r0 res'+                | otherwise = res'+           in res''++-- | \(O(\min(|S|, L) \log |S|)\) Collects \((y, \mathrm{rank}(y))\) in range \([l, r)\) in+-- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small.+--+-- @since 1.1.0.0+{-# INLINE descAssocsIn #-}+descAssocsIn :: RawWaveletMatrix -> Int -> Int -> [(Int, Int)]+descAssocsIn wm l r = descAssocsInWith wm l r id++-- | \(O(\log A \min(|A|, L))\) Internal implementation of `descAssoc`.+--+-- @since 1.1.0.0+{-# INLINE descAssocsInWith #-}+descAssocsInWith :: RawWaveletMatrix -> Int -> Int -> (Int -> Int) -> [(Int, Int)]+descAssocsInWith RawWaveletMatrix {..} l_ r_ f+  | l'_ < r'_ = inner (0 :: Int) (0 :: Int) l'_ r'_ []+  | otherwise = []+  where+    -- clamp+    l'_ = max 0 l_+    r'_ = min lengthRwm r_+    -- DFS. [l, r)+    inner !acc iRow !l !r res+      | iRow >= heightRwm =+          let !n = r - l+              !acc' = f acc+           in (acc', n) : res+      | otherwise = do+          let !bits = bitsRwm VG.! iRow+              !l0 = BV.rank0 bits l+              !r0 = BV.rank0 bits r+              !nZeros = nZerosRwm VG.! iRow+              !res'+                -- go left+                | l0 < r0 = inner acc (iRow + 1) l0 r0 res+                | otherwise = res+              -- go right (visit bigger values first)+              !l' = l + nZeros - l0+              !r' = r + nZeros - r0+              !res''+                | l' < r' = inner (acc .|. bit (heightRwm - 1 - iRow)) (iRow + 1) l' r' res'+                | otherwise = res'+           in res''

@@ -0,0 +1,277 @@+{-# LANGUAGE RecordWildCards #-}++-- | A 2D, static wavelet matrix with segment tree, that can handle point add and rectangle sum+-- queries. Points cannot be added after construction, but monoid values in each point can be+-- modified later.+--+-- ==== __Example__+-- Create a `WaveletMatrix2d` with initial vertex values:+--+-- >>> import AtCoder.Extra.WaveletMatrix2d qualified as WM+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> -- 8  9 10 11+-- >>> -- 4  5  6  7+-- >>> -- 0  1  2  3+-- >>> wm <- WM.build negate $ VU.generate 12 $ \i -> let (!y, !x) = i `divMod` 4 in (x, y, Sum i)+--+-- Read the value at \(x = 2, y = 1\):+--+-- >>> WM.read wm (2, 1)+-- Sum {getSum = 6}+--+-- Other segment tree methods are also available, but in 2D:+--+-- >>> WM.allProd wm -- (0 + 11) * 12 / 2 = 66+-- Sum {getSum = 66}+--+-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 5 + 6 + 9 + 10+-- Sum {getSum = 33}+--+-- >>> WM.modify wm (+ 2) (1, 1)+-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 7 + 6 + 9 + 10+-- Sum {getSum = 35}+--+-- >>> WM.write wm (1, 1) $ Sum 0+-- >>> WM.prod wm {- x -} 1 3 {- y -} 0 3 -- 1 + 2 + 0 + 6 + 9 + 10+-- Sum {getSum = 28}+module AtCoder.Extra.WaveletMatrix2d+  ( -- * Wavelet matrix 2D+    WaveletMatrix2d (..),++    -- * Counstructor+    new,+    build,++    -- * Segment tree methods+    read,+    write,+    modify,+    prod,+    prodMaybe,+    allProd,+    -- wavelet matrix methods could be implemented, too+  )+where++import AtCoder.Extra.Bisect (bisectR, lowerBound)+import AtCoder.Extra.WaveletMatrix.BitVector qualified as BV+import AtCoder.Extra.WaveletMatrix.Raw qualified as Rwm+import AtCoder.Internal.Assert qualified as ACIA+import AtCoder.SegTree qualified as ST+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Bit (Bit (..))+import Data.Bits (Bits (testBit))+import Data.Maybe (fromJust, fromMaybe)+import Data.Vector qualified as V+import Data.Vector.Algorithms.Intro qualified as VAI+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU+import GHC.Stack (HasCallStack)+import Prelude hiding (read)++-- NOTE: There are many possible improvements.+-- - Use cumulative sum or fenwick tree instead for the speed.+-- - The inverse operator is not actually required.+-- - Wavelet matrix methods such as `rank` can be implemented+-- - `maxRight` can be implemented.++-- | Segment Tree on Wavelet Matrix: points on a 2D plane and rectangle products.+data WaveletMatrix2d s a = WaveletMatrix2d+  { -- | The wavelet matrix that represents points on a 2D plane.+    rawWmWm2d :: !Rwm.RawWaveletMatrix,+    -- | (x, y) index compression dictionary.+    xyDictWm2d :: !(VU.Vector (Int, Int)),+    -- | y index compression dictionary.+    yDictWm2d :: !(VU.Vector Int),+    -- | The segment tree of the weights of the points in the order of `xyDictWm2d`.+    segTreesWm2d :: !(V.Vector (ST.SegTree s a)),+    -- | The inverse operator of the interested monoid.+    invWm2d :: !(a -> a)+  }++-- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with `mempty` as the initial monoid+-- values for each point.+{-# INLINE new #-}+new ::+  (PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Inverse operator of the monoid+  (a -> a) ->+  -- | Input points+  VU.Vector (Int, Int) ->+  -- | A 2D wavelet matrix+  m (WaveletMatrix2d (PrimState m) a)+new invWm2d xys = do+  let n = VG.length xys+  let xyDictWm2d = VU.uniq . VU.modify (VAI.sortBy compare) $ xys+  let (!_, !ys) = VU.unzip xys+  let yDictWm2d = VU.uniq $ VU.modify (VAI.sortBy compare) ys+  -- REMARK: Be sure to use `n + 1` because the product function cannot handle the case+  --         `yUpper` is `2^{height}`.+  let (!_, !ysInput) = VU.unzip xyDictWm2d+  let rawWmWm2d = Rwm.build (n + 1) $ VU.map (fromJust . lowerBound yDictWm2d) ysInput+  segTreesWm2d <- V.replicateM (Rwm.heightRwm rawWmWm2d) (ST.new n)+  pure WaveletMatrix2d {..}++-- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with wavelet matrix with segment tree+-- with initial monoid values. Monoids on a duplicate point are accumulated with `(<>)`.+{-# INLINE build #-}+build ::+  (PrimMonad m, Monoid a, VU.Unbox a) =>+  -- | Inverse operator of the monoid+  (a -> a) ->+  -- | Input points with initial values+  VU.Vector (Int, Int, a) ->+  -- | A 2D wavelet matrix+  m (WaveletMatrix2d (PrimState m) a)+build invWm2d xysw = do+  let (!xs, !ys, !_) = VU.unzip3 xysw+  wm <- new invWm2d $ VU.zip xs ys+  -- not the fastest implementation though+  VU.forM_ xysw $ \(!x, !y, !w) -> do+    modify wm (<> w) (x, y)+  pure wm++-- | \(O(1)\) Returns the monoid value at \((x, y)\).+{-# INLINE read #-}+read :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (Int, Int) -> m a+read WaveletMatrix2d {..} (!x, !y) = do+  ST.read (V.head segTreesWm2d) . fromJust $ lowerBound xyDictWm2d (x, y)++-- | \(O(\log^2 n)\) Writes the monoid value at \((x, y)\). Access to unknown points are undefined.+{-# INLINE write #-}+write :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (Int, Int) -> a -> m ()+write WaveletMatrix2d {..} (!x, !y) v = do+  let !i_ = fromJust $ lowerBound xyDictWm2d (x, y)+  V.ifoldM'_+    ( \i iRow (!bits, !seg) -> do+        let !i0 = BV.rank0 bits i+        let !i'+              | unBit $ VG.unsafeIndex (BV.bitsBv bits) i =+                  i + Rwm.nZerosRwm rawWmWm2d VG.! iRow - i0+              | otherwise = i0+        ST.write seg i' v+        pure i'+    )+    i_+    $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d++-- | \(O(\log^2 n)\) Modifies the monoid value at \((x, y)\). Access to unknown points are+-- undefined.+{-# INLINE modify #-}+modify :: (HasCallStack, Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> (a -> a) -> (Int, Int) -> m ()+modify WaveletMatrix2d {..} f (!x, !y) = do+  let !i_ = fromJust $ lowerBound xyDictWm2d (x, y)+  V.ifoldM'_+    ( \i iRow (!bits, !seg) -> do+        let !i0 = BV.rank0 bits i+        let !i'+              | unBit $ VG.unsafeIndex (BV.bitsBv bits) i =+                  i + Rwm.nZerosRwm rawWmWm2d VG.! iRow - i0+              | otherwise = i0+        ST.modify seg f i'+        pure i'+    )+    i_+    $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d++-- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\).+{-# INLINE prod #-}+prod :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a+prod wm@WaveletMatrix2d {..} !xl !xr !yl !yr+  | xl' >= xr' || yl' >= yr' = pure mempty+  | otherwise = unsafeProd wm xl' xr' yl' yr'+  where+    (!xDict, !_) = VU.unzip xyDictWm2d+    -- NOTE: clamping here!+    xl' = fromMaybe 0 $ bisectR 0 (VG.length xDict) $ (< xl) . VG.unsafeIndex xDict+    xr' = fromMaybe (VG.length xDict) $ bisectR 0 (VG.length xDict) $ (< xr) . VG.unsafeIndex xDict+    yl' = fromMaybe 0 $ bisectR 0 (VG.length yDictWm2d) $ (< yl) . VG.unsafeIndex yDictWm2d+    yr' = fromMaybe (VG.length yDictWm2d) $ bisectR 0 (VG.length yDictWm2d) $ (< yr) . VG.unsafeIndex yDictWm2d+    !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed x)" xl' xr' (VG.length xDict)+    !_ = ACIA.checkInterval "AtCoder.Extra.WaveletMatrix.SegTree.prod (compressed y)" yl' yr' (VG.length yDictWm2d)++-- | \(O(\log^2 n)\) Returns the monoid product in \([l, r) \times [y_1, y_2)\). Returns `Nothing` for invalid+-- intervals.+{-# INLINE prodMaybe #-}+prodMaybe :: (VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m (Maybe a)+prodMaybe wm@WaveletMatrix2d {..} !xl !xr !yl !yr+  | not (ACIA.testInterval xl' xr' (VG.length xDict)) = pure Nothing+  | not (ACIA.testInterval yl' yr' (VG.length yDictWm2d)) = pure Nothing+  | xl' >= xr' || yl' >= yr' = pure $ Just mempty+  | otherwise = Just <$> unsafeProd wm xl' xr' yl' yr'+  where+    (!xDict, !_) = VU.unzip xyDictWm2d+    -- NOTE: clamping here!+    xl' = fromMaybe 0 $ bisectR 0 (VG.length xDict) $ (< xl) . VG.unsafeIndex xDict+    xr' = fromMaybe (VG.length xDict) $ bisectR 0 (VG.length xDict) $ (< xr) . VG.unsafeIndex xDict+    yl' = fromMaybe 0 $ bisectR 0 (VG.length yDictWm2d) $ (< yl) . VG.unsafeIndex yDictWm2d+    yr' = fromMaybe (VG.length yDictWm2d) $ bisectR 0 (VG.length yDictWm2d) $ (< yr) . VG.unsafeIndex yDictWm2d++-- | \(O(\log^2 n)\) Return the monoid product of all of the points in the wavelet matrix.+{-# INLINE allProd #-}+allProd :: (HasCallStack, VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> m a+allProd WaveletMatrix2d {..} = do+  -- ST.allProd (V.last segTreesWm2d)+  ST.allProd (V.head segTreesWm2d)++-- | \(O(\log^2 n)\) The input is compressed indices.+{-# INLINE unsafeProd #-}+unsafeProd :: (VU.Unbox a, Monoid a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a+unsafeProd wm xl' xr' yl' yr' = do+  sR <- prodLT wm xl' xr' yr'+  sL <- prodLT wm xl' xr' yl'+  pure $! sR <> invWm2d wm sL++-- | \(O(\log^2 n)\)+{-# INLINE prodLT #-}+prodLT :: (Monoid a, VU.Unbox a, PrimMonad m) => WaveletMatrix2d (PrimState m) a -> Int -> Int -> Int -> m a+prodLT WaveletMatrix2d {..} !l_ !r_ yUpper = do+  (!res, !_, !_) <- do+    V.ifoldM'+      ( \(!acc, !l, !r) !iRow (!bits, !seg) -> do+          let !l0 = BV.rank0 bits l+              !r0 = BV.rank0 bits r+          -- REMARK: The function cannot handle the case yUpper = N = 2^i. See the constructor for+          -- how it's handled and note that l_ and r_ are compressed indices.+          if testBit yUpper (Rwm.heightRwm rawWmWm2d - 1 - iRow)+            then do+              !acc' <- (acc <>) <$> ST.prod seg l0 r0+              let !l' = l + Rwm.nZerosRwm rawWmWm2d VG.! iRow - l0+              let !r' = r + Rwm.nZerosRwm rawWmWm2d VG.! iRow - r0+              pure (acc', l', r')+            else do+              pure (acc, l0, r0)+      )+      (mempty, l_, r_)+      $ V.zip (Rwm.bitsRwm rawWmWm2d) segTreesWm2d+  pure res++-- -- | \(O(\log n)\) Restore the original \(x\) coordinate from a compressed one. Access to unknown+-- -- points are undefined.+-- {-# INLINE indexX #-}+-- indexX :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int+-- indexX WaveletMatrix2d {xyDictWm2d} x = maybe err (VG.unsafeIndex xDict) $ lowerBound xDict x+--   where+--     (!xDict, !_) = VU.unzip xyDictWm2d+--     err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexX: cannot index x (`" ++ show x ++ "`)"++-- -- | \(O(\log n)\) Restore the original \(y\) coordinate from a compressed one. Access to unknown+-- -- points are undefined.+-- {-# INLINE indexY #-}+-- indexY :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int+-- indexY WaveletMatrix2d {yDictWm2d} y = maybe err (VG.unsafeIndex yDictWm2d) $ lowerBound yDictWm2d y+--   where+--     err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexY: cannot index y (`" ++ show y ++ "`)"++-- -- | \(O(\log n)\) Restore the original \((x, y)\) coordinates from a compressed one. Access to+-- -- unknown points are undefined.+-- {-# INLINE indexXY #-}+-- indexXY :: (HasCallStack) => WaveletMatrix2d s a -> Int -> Int -> (Int, Int)+-- indexXY WaveletMatrix2d {xyDictWm2d} x y = maybe err (VG.unsafeIndex xyDictWm2d) $ lowerBound xyDictWm2d (x, y)+--   where+--     err = error $ "AtCoder.Extra.WaveletMatirx.SegTree.indexXY: cannot index (x, y) `" ++ show (x, y) ++ "`"++-- {-# INLINE assocsWith #-}+-- assocsWith :: WaveletMatrix -> (Int -> Int) -> [(Int, Int)]+-- assocsWith WaveletMatrix {..} l_ r_ f