ac-library-hs 1.5.2.1 → 1.5.3.0

raw patch · 9 files changed

+949/−14 lines, 9 filesdep ~randomdep ~unordered-containersPVP ok

version bump matches the API change (PVP)

Dependency ranges changed: random, unordered-containers

API changes (from Hackage documentation)

Files

• CHANGELOG.md +6/−0
• ac-library-hs.cabal +16/−10
• benchmarks/Bench/MulMod.hs +64/−0
• src/AtCoder/Extra/AhoCorasick.hs +284/−0
• src/AtCoder/Extra/DsuMonoid.hs +248/−0
• src/AtCoder/Extra/KdTree.hs +3/−3
• test/Main.hs +5/−1
• test/Tests/Extra/Vector.hs +171/−0
• test/Tests/Extra/Vector/Prim.hs +152/−0

@@ -1,5 +1,11 @@ # Revision history for acl-hs +## 1.5.3.0 -- June 2025++- Added `AtCoder.Extra.AhoCorasick`+- Added `AtCoder.Extra.DsuMonoid`+- Fixed the implementation of `KdTree.findNearstPoint`+ ## 1.5.2.0 -- June 2025  - Added `AtCoder.Extra.Vector.Prim`

@@ -4,7 +4,7 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.5.2.1+version:         1.5.3.0 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive@@ -35,15 +35,16 @@  common dependencies   build-depends:-    , base               >=4.9     && <4.22-    , bitvec             <1.2-    , bytestring         <0.14-    , primitive          >=0.6.4.0 && <0.10-    , random             >=1.2.0   && <1.3-    , transformers       >=0.2.0.0-    , vector             >=0.13.0  && <0.14-    , vector-algorithms  <0.10-    , wide-word          <0.2+    , base                  >=4.9     && <4.22+    , bitvec                <1.2+    , bytestring            <0.14+    , primitive             >=0.6.4.0 && <0.10+    , random                >=1.2.0   && <1.4+    , transformers          >=0.2.0.0+    , unordered-containers  <0.3+    , vector                >=0.13.0  && <0.14+    , vector-algorithms     <0.10+    , wide-word             <0.2    default-language: GHC2021 @@ -53,7 +54,9 @@   exposed-modules:     AtCoder.Convolution     AtCoder.Dsu+    AtCoder.Extra.AhoCorasick     AtCoder.Extra.Bisect+    AtCoder.Extra.DsuMonoid     AtCoder.Extra.DynLazySegTree     AtCoder.Extra.DynLazySegTree.Persistent     AtCoder.Extra.DynLazySegTree.Raw@@ -174,6 +177,8 @@     Tests.Extra.Seq.Map     Tests.Extra.Tree     Tests.Extra.Tree.Lct+    Tests.Extra.Vector+    Tests.Extra.Vector.Prim     Tests.Extra.WaveletMatrix     Tests.Extra.WaveletMatrix.BitVector     Tests.Extra.WaveletMatrix.Raw@@ -230,6 +235,7 @@     Bench.Matrix     Bench.ModInt     Bench.Montgomery64+    Bench.MulMod     Bench.PowMod     Bench.RepeatWithIndex     Bench.RepeatWithoutIndex

@@ -0,0 +1,64 @@+module Bench.MulMod (benches) where++import BenchLib.MulMod.Barrett64 qualified as Barrett64+import BenchLib.MulMod.BarrettWideWord qualified as BarrettWideWord+import BenchLib.MulMod.Montgomery qualified as Montgomery+import Criterion+import Data.Bifunctor (first)+import Data.Vector.Unboxed qualified as VU+import Data.Word (Word32, Word64)+import System.Random++n :: Int+n = 10000++btWW :: BarrettWideWord.Barrett+btWW = BarrettWideWord.new64 998244353++bt64 :: Barrett64.Barrett+bt64 = Barrett64.new 998244353++mont :: Montgomery.Montgomery+mont = Montgomery.new 998244353++-- | This benchmark is almost nonsense, as the pre calculation' overhead is not considered while the+-- real use case is @powMod@. However, it's useful when optimizing the successive calculation.+benches32 :: Benchmark+benches32 =+  bgroup+    "mulMod Word32 vector"+    [ bench "barrettWideWord" $ whnf (VU.foldl' (\acc -> BarrettWideWord.mulMod btWW acc . fromIntegral) w64) nonZeroRandomVec32,+      bench "barrett" $ whnf (VU.foldl' (\acc -> Barrett64.mulMod bt64 acc . fromIntegral) w64) nonZeroRandomVec32,+      -- NOTE: It skips the last reduce.+      bench "montgomery" $ whnf (VU.foldl' (\acc -> Montgomery.mulModGenerated mont acc . Montgomery.generate mont . fromIntegral) w64) nonZeroRandomVec32,+      bench "mod" $ whnf (VU.foldl' (\acc n -> acc * n `mod` 998244353) w32) nonZeroRandomVec32,+      bench "rem" $ whnf (VU.foldl' (\acc n -> acc * n `rem` 998244353) w32) nonZeroRandomVec32+    ]+  where+    w32 :: Word32+    w32 = 1+    w64 :: Word64+    w64 = 1+    -- [1, 998244383)+    nonZeroRandomVec32 :: VU.Vector Word32+    nonZeroRandomVec32 = VU.map fromIntegral $ VU.unfoldrExactN n (first (+ 1) . genWord64R (998244383 - 2)) (mkStdGen 123456789)++-- | This benchmark is almost nonsense, as the pre calculation' overhead is not considered while the+-- real use case is @powMod@. However, it's useful when optimizing the successive calculation.+benches :: Benchmark+benches =+  bgroup+    "mulMod Word64 vector"+    [ bench "barrettWideWord" $ whnf (VU.foldl' (BarrettWideWord.mulMod btWW) w64) nonZeroRandomVec64,+      bench "barrett" $ whnf (VU.foldl' (\acc -> Barrett64.mulMod bt64 acc) w64) nonZeroRandomVec64,+      -- NOTE: It skips the last reduce.+      bench "montgomery" $ whnf (VU.foldl' (\acc -> Montgomery.mulModGenerated mont acc . Montgomery.generate mont) w64) nonZeroRandomVec64,+      bench "mod" $ whnf (VU.foldl' (\acc n -> acc * n `mod` 998244353) w64) nonZeroRandomVec64,+      bench "rem" $ whnf (VU.foldl' (\acc n -> acc * n `rem` 998244353) w64) nonZeroRandomVec64+    ]+  where+    w64 :: Word64+    w64 = 1+    -- [1, 998244383)+    nonZeroRandomVec64 :: VU.Vector Word64+    nonZeroRandomVec64 = VU.map fromIntegral $ VU.unfoldrExactN n (first (+ 1) . genWord64R (998244383 - 2)) (mkStdGen 123456789)

@@ -0,0 +1,284 @@+-- | Aho–Corasick algorithm is a fast dictionary-matching (multi-pattern matching) algorithm.+--+-- ==== __Example__+--+-- >>> import AtCoder.Extra.AhoCorasick qualified as AC+-- >>> import Data.Vector.Unboxed qualified as VU+--+-- Pattern strings must be given as @V.Vector (VU.Vector Int)@:+--+-- >>> let patterns = V.fromList [VU.fromList [0, 1], VU.fromList [0, 2], VU.fromList [2, 3, 4]]+-- >>> let ac = AC.build patterns+-- >>> AC.size ac+-- 7+--+-- The automaton could be run manually with `next` or `nextN`:+--+-- >>> AC.nextN ac {- empty node -} 0 (VU.fromList [0, 2, 3])+-- 5+--+-- `match` returns a vector of @(endPos, patternId)@:+--+-- >>> --                         [.....) pattern 0+-- >>> --                               [.......) pattern2+-- >>> AC.match ac $ VU.fromList [0, 1, 2, 3, 4]+-- [(2,0),(5,2)]+--+-- If you need a vector of @(startPos, patternId)@, you must manually map the result:+--+-- >>> let f (!end, !patId) = (end - VU.length (patterns V.! patId), patId)+-- >>> --                                    [.....) pattern 0+-- >>> --                                          [.......) pattern2+-- >>> VU.map f . AC.match ac $ VU.fromList [0, 1, 2, 3, 4]+-- [(0,0),(2,2)]+--+-- Note that duplicate patterns are only counted once with `match`.+--+-- @since 1.5.3.0+module AtCoder.Extra.AhoCorasick+  ( AhoCorasick (..),+    build,+    size,+    next,+    nextN,+    match,+  )+where++-- TODO: Generalize with Hash + Unbox? Int-only implementation is faster though.++import AtCoder.Extra.Vector qualified as EV+import AtCoder.Internal.Queue qualified as Q+import Control.Monad (when)+import Control.Monad.Fix (fix)+import Control.Monad.ST (runST)+import Data.Foldable (for_)+import Data.HashMap.Strict qualified as HM+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Generic.Mutable qualified as VGM+import Data.Vector.Mutable qualified as VM+import Data.Vector.Unboxed qualified as VU+import Data.Vector.Unboxed.Mutable qualified as VUM+import GHC.Stack (HasCallStack)++-- | Aho–Corasick algorithm data.+--+-- @since 1.5.3.0+data AhoCorasick = AhoCorasick+  { -- | The number of nodes in the trie.+    --+    -- @since 1.5.3.0+    sizeAc :: {-# UNPACK #-} !Int,+    -- | A trie (-like directed graph) of input words: Vertex -> (Char -> Vertex).+    --+    -- @since 1.5.3.0+    trieAc :: !(V.Vector (HM.HashMap Int Int)),+    -- | Node data of links to parent vertex.+    --+    -- @since 1.5.3.0+    parentAc :: !(VU.Vector Int),+    -- | Node data that represents completed pattern string or nothing (@-1@).+    --+    -- @since 1.5.3.0+    patternAc :: !(VU.Vector Int),+    -- | Node data of links to the longest suffix vertex.+    --+    -- @since 1.5.3.0+    suffixAc :: !(VU.Vector Int),+    -- | Node data of links to the longest suffix pattern vertex.+    --+    -- @since 1.5.3.0+    outputAc :: !(VU.Vector Int)+  }++-- | \(O(\sum_i |S_i|)\)+--+-- ==== Constraints+-- - \(|S_i| > 0\)+--+-- @since 1.5.3.0+{-# INLINEABLE build #-}+build ::+  (HasCallStack) =>+  -- | Pattern strings.+  V.Vector (VU.Vector Int) ->+  -- | Aho–Corasick automaton based on a trie.+  AhoCorasick+build patterns+  | VG.null patterns =+      -- root only+      AhoCorasick+        1+        (V.singleton HM.empty)+        (VU.replicate 1 (-1))+        (VU.replicate 1 0)+        (VU.replicate 1 0)+        (VU.replicate 1 0)+  | otherwise =+      let (!nNodes, !patternMap, !trie, !parent) = buildTrie patterns+          (!suffix, !output) = runBfs nNodes trie patternMap+       in AhoCorasick nNodes trie parent patternMap suffix output++-- | \(O(1)\) Returns the number of nodes in the trie.+--+-- @since 1.5.3.0+{-# INLINE size #-}+size :: (HasCallStack) => AhoCorasick -> Int+size = sizeAc++-- | \(O(1)\) Retrieves the next node to visit.+--+-- @since 1.5.3.0+{-# INLINEABLE next #-}+-- TODO: benchmark INLINE+next ::+  (HasCallStack) =>+  -- | The automaton.+  AhoCorasick ->+  -- | Current node ID (empty node is @0@).+  Int ->+  -- | Character.+  Int ->+  -- | Next node ID.+  Int+next AhoCorasick {trieAc, suffixAc} v0 c = inner v0+  where+    inner v = case HM.lookup c (trieAc VG.! v) of+      Just end -> end+      Nothing+        -- no hope+        | v == 0 -> 0+        -- fallback to the longest match suffix+        | otherwise -> inner $! suffixAc VG.! v++-- | \(n\) Applies `next` N times for a given input string.+--+-- ==== Constraints+--+-- @since 1.5.3.0+{-# INLINE nextN #-}+nextN ::+  (HasCallStack) =>+  -- | The automaton.+  AhoCorasick ->+  -- | Current node.+  Int ->+  -- | String.+  VU.Vector Int ->+  -- | Resulting node.+  Int+nextN ac = VU.foldl' (next ac)++-- | \(O(|T|)\) Runs dictionary matching (multi-pattern matching) in linear time and returns a list+-- of @(endPos, patId)@, where @[endPos - patLen, endPos)@ corresponds to the interval of original+-- source slice.+--+-- Note that duplicate patterns are counted just once with one of them; if pattern A and B are the+-- same, their appearence is counted as either A or B.+--+-- @since 1.5.3.0+{-# INLINEABLE match #-}+match :: (HasCallStack) => AhoCorasick -> VU.Vector Int -> VU.Vector (Int, Int)+match ac@AhoCorasick {patternAc, outputAc} =+  EV.iconcatMap (\i v -> VU.unfoldr (f i) v) . VU.postscanl' (next ac) 0+  where+    f :: Int -> Int -> Maybe ((Int, Int), Int)+    f _ 0 = Nothing+    f i v = case patternAc VG.! v of+      -- NOTE: Do not perform early return, as the initial vertex can be non-pattern+      -1 -> f i (outputAc VG.! v)+      -- NOTE: Here we use `i + 1`, where [pos - patLen, pos) makes up a half-open interval.+      pat -> Just ((i + 1, pat), outputAc VG.! v)++-- | \(O(\sum_i |S_i| \Gamma)\)+{-# INLINEABLE buildTrie #-}+buildTrie :: (HasCallStack) => V.Vector (VU.Vector Int) -> (Int, VU.Vector Int, V.Vector (HM.HashMap Int Int), VU.Vector Int)+buildTrie patternStrings = runST $ do+  let !nMaxNodes = (1 +) . V.sum $ V.map VG.length patternStrings++  -- allocator+  nNodesVec <- VUM.replicate 1 (1 :: Int)++  -- components+  nextVec <- VM.replicate nMaxNodes HM.empty+  parentVec <- VUM.replicate nMaxNodes (0 :: Int)++  -- create a trie and collect pattern vertices+  patternVerts <-+    (VU.convert <$>) . V.forM patternStrings $+      VG.foldM'+        ( \ !u c -> do+            v0 <- HM.lookup c <$> VGM.read nextVec u+            case v0 of+              Nothing -> do+                -- allocate a new vertex index+                v <- VGM.read nNodesVec 0+                VGM.write nNodesVec 0 $! v + 1+                -- store the next vertex link+                VGM.modify nextVec (HM.insert c v) u+                -- fill the vertex information+                VGM.write parentVec v u+                pure v+              Just v -> do+                pure v+        )+        0++  !nNodes <- VGM.read nNodesVec 0++  let !pattern = VU.create $ do+        -- We could replace the following with VU.accumulate+        patVec <- VUM.replicate nNodes (-1 :: Int)+        VU.iforM_ patternVerts $ \iPattern v -> do+          VGM.write patVec v iPattern+        pure patVec+  !trie <- VG.take nNodes <$> V.unsafeFreeze nextVec+  !parent <- VG.take nNodes <$> VU.unsafeFreeze parentVec+  pure (nNodes, pattern, trie, parent)++-- | \(O(\sum_i |S_i| \Gamma)\)+{-# INLINEABLE runBfs #-}+runBfs :: (HasCallStack) => Int -> V.Vector (HM.HashMap Int Int) -> VU.Vector Int -> (VU.Vector Int, VU.Vector Int)+runBfs nNodes trie patternMap = runST $ do+  suffixVec <- VUM.replicate nNodes (0 :: Int)+  outputVec <- VUM.replicate nNodes (0 :: Int)++  que <- Q.new @_ @Int nNodes+  for_ (HM.elems (trie VG.! 0)) $ \v -> do+    Q.pushBack que v++  -- TODO: deduplicate with `next`+  let nextM c u = case HM.lookup c (trie VG.! u) of+        Just end -> pure end+        Nothing+          | u == 0 -> pure 0+          | otherwise -> do+              v <- VGM.read suffixVec u+              nextM c v++  fix $ \popLoop -> do+    q <- Q.popFront que+    case q of+      Nothing -> pure ()+      Just u -> do+        -- visit neighbors+        for_ (HM.toList (trie VG.! u)) $ \(!c, !v) -> do+          Q.pushBack que v++          -- find the longest suffix to continue with `c`+          !suffix <- nextM c =<< VGM.read suffixVec u+          VGM.write suffixVec v suffix++          -- find the longest suffix that matches to a pattern+          let suffixPattern = patternMap VG.! suffix+          output <-+            if suffixPattern /= -1+              then pure suffix+              else VGM.read outputVec suffix+          VGM.write outputVec v output++        -- loop+        popLoop++  (,) <$> VU.unsafeFreeze suffixVec <*> VU.unsafeFreeze outputVec

@@ -0,0 +1,248 @@+{-# LANGUAGE RecordWildCards #-}++-- | A disjoint set union with commutative monoid values associated with each group.+--+-- ==== __Example__+--+-- >>> import AtCoder.Extra.DsuMonoid qualified as Dm+-- >>> import Data.Semigroup (Sum (..))+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> dsu <- Dm.build $ VU.generate 4 Sum+-- >>> Dm.merge dsu 0 1+-- 0+--+-- >>> Dm.read dsu 0+-- Sum {getSum = 1}+--+-- >>> Dm.read dsu 1+-- Sum {getSum = 1}+--+-- >>> Dm.mergeMaybe dsu 0 2+-- Just 0+--+-- >>> Dm.read dsu 0+-- Sum {getSum = 3}+--+-- @since 1.5.3.0+module AtCoder.Extra.DsuMonoid+  ( -- * Disjoint set union+    DsuMonoid (dsuDm, mDm),++    -- * Constructors+    new,+    build,++    -- * Merging+    merge,+    mergeMaybe,+    merge_,++    -- * Leader+    leader,++    -- * Component information+    same,+    size,+    groups,++    -- * Monoid values+    read,+    unsafeRead,+    unsafeWrite,+  )+where++import AtCoder.Dsu qualified as Dsu+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+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+import GHC.Stack (HasCallStack)+import Prelude hiding (read)++-- | A disjoint set union with commutative monoid values associated with each group.+--+-- @since 1.5.3.0+data DsuMonoid s a = DsuMonoid+  { -- | The original DSU.+    --+    -- @since 1.5.3.0+    dsuDm :: {-# UNPACK #-} !(Dsu.Dsu s),+    -- | Commutative monoid values for each group.+    --+    -- @since 1.5.3.0+    mDm :: !(VUM.MVector s a)+  }++-- | Creates an undirected graph with \(n\) vertices and \(0\) edges.+--+-- ==== Constraints+-- - \(0 \le n\)+--+-- ==== Complexity+-- - \(O(n)\)+--+-- @since 1.5.3.0+{-# INLINE new #-}+new :: (PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (DsuMonoid (PrimState m) a)+new n+  | n >= 0 = build $ VU.replicate n mempty+  | otherwise = error $ "AtCoder.Extra.DsuMonoid: given negative size (`" ++ show n ++ "`)"++-- | Creates an undirected graph with \(n\) vertices and \(0\) edges.+--+-- ==== Constraints+-- - \(0 \le n\)+--+-- ==== Complexity+-- - \(O(n)\)+--+-- @since 1.5.3.0+{-# INLINE build #-}+build :: (PrimMonad m, Semigroup a, VU.Unbox a) => VU.Vector a -> m (DsuMonoid (PrimState m) a)+build ms = stToPrim $ do+  dsuDm <- Dsu.new $ VU.length ms+  mDm <- VU.thaw ms+  pure $ DsuMonoid {..}++-- | 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+-- - \(0 \leq a < n\)+-- - \(0 \leq b < n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\) amortized+--+-- @since 1.5.3.0+{-# INLINEABLE merge #-}+merge :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> Int -> m Int+merge DsuMonoid {..} a b = stToPrim $ do+  r1 <- Dsu.leader dsuDm a+  r2 <- Dsu.leader dsuDm b+  if r1 == r2+    then pure r1+    else do+      !m1 <- VGM.read mDm r1+      !m2 <- VGM.read mDm r2+      r' <- Dsu.merge dsuDm a b+      VGM.write mDm r' $! m1 <> m2+      pure r'++-- | Adds an edge \((a, b)\). It returns the representative of the new connected component, or+-- `Nothing` if the two vertices are in the same connected component.+--+-- ==== Constraints+-- - \(0 \leq a < n\)+-- - \(0 \leq b < n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\) amortized+--+-- @since 1.2.4.0+{-# INLINEABLE mergeMaybe #-}+mergeMaybe :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> Int -> m (Maybe Int)+mergeMaybe DsuMonoid {..} a b = stToPrim $ do+  r1 <- Dsu.leader dsuDm a+  r2 <- Dsu.leader dsuDm b+  if r1 == r2+    then pure Nothing+    else do+      !m1 <- VGM.read mDm r1+      !m2 <- VGM.read mDm r2+      r' <- Dsu.merge dsuDm a b+      VGM.write mDm r' $! m1 <> m2+      pure $ Just r'++-- | `merge` with the return value discarded.+--+-- ==== Constraints+-- - \(0 \leq a < n\)+-- - \(0 \leq b < n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\) amortized+--+-- @since 1.5.3.0+{-# INLINE merge_ #-}+merge_ :: (PrimMonad m, Semigroup a, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> Int -> m ()+merge_ dsu a b = do+  _ <- merge dsu a b+  pure ()++-- | Returns whether the vertices \(a\) and \(b\) are in the same connected component.+--+-- ==== Constraints+-- - \(0 \leq a < n\)+-- - \(0 \leq b < n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\) amortized+--+-- @since 1.5.3.0+{-# INLINE same #-}+same :: (HasCallStack, PrimMonad m) => DsuMonoid (PrimState m) a -> Int -> Int -> m Bool+same dsu = Dsu.same (dsuDm dsu)++-- | Returns the representative of the connected component that contains the vertex \(a\).+--+-- ==== Constraints+-- - \(0 \leq a \lt n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\) amortized+--+-- @since 1.5.3.0+{-# INLINE leader #-}+leader :: (HasCallStack, PrimMonad m) => DsuMonoid (PrimState m) a -> Int -> m Int+leader dsu = Dsu.leader (dsuDm dsu)++-- | Returns the size of the connected component that contains the vertex \(a\).+--+-- ==== Constraints+-- -  \(0 \leq a < n\)+--+-- ==== Complexity+-- - \(O(\alpha(n))\)+--+-- @since 1.5.3.0+{-# INLINE size #-}+size :: (HasCallStack, PrimMonad m) => DsuMonoid (PrimState m) a -> Int -> m Int+size dsu = Dsu.size (dsuDm dsu)++-- | \O(n)\) 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.+--+-- @since 1.5.3.0+{-# INLINE groups #-}+groups :: (PrimMonad m) => DsuMonoid (PrimState m) a -> m (V.Vector (VU.Vector Int))+groups dsu = Dsu.groups (dsuDm dsu)++-- | \(O(1)\) Reads the group value of the \(k\)-th node.+--+-- @since 1.5.3.0+{-# INLINE read #-}+read :: (PrimMonad m, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> m a+read DsuMonoid {..} i = do+  VGM.read mDm =<< Dsu.leader dsuDm i++-- | \(O(1)\) Reads the group value of the \(k\)-th node.+--+-- @since 1.5.3.0+{-# INLINE unsafeRead #-}+unsafeRead :: (PrimMonad m, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> m a+unsafeRead DsuMonoid {..} i = do+  VGM.read mDm i++-- | \(O(1)\) Writes to the group value of the \(k\)-th node.+--+-- @since 1.5.3.0+{-# INLINE unsafeWrite #-}+unsafeWrite :: (PrimMonad m, VU.Unbox a) => DsuMonoid (PrimState m) a -> Int -> a -> m ()+unsafeWrite DsuMonoid {..} i x = do+  VGM.write mDm i x

@@ -226,10 +226,10 @@       -- it's a leaf       | dataI /= -1 = (dataI, d)       -- look into the children-      | d0 < d1 = inner (2 * i + 0) $ inner (2 * i + 1) res-      | otherwise = inner (2 * i + 1) $ inner (2 * i + 0) res+      | d0 < d1 = inner (2 * i + 1) $ inner (2 * i + 0) res+      | otherwise = inner (2 * i + 0) $ inner (2 * i + 1) res       where         d = bestDistSquared i         dataI = dataKt VG.! i         d0 = bestDistSquared (2 * i + 0)-        d1 = bestDistSquared (2 * i + 0)+        d1 = bestDistSquared (2 * i + 1)

@@ -4,7 +4,9 @@ import Test.Tasty.Ingredients.Rerun import Tests.Convolution qualified import Tests.Dsu qualified+import Tests.Extra.AhoCorasick qualified import Tests.Extra.Bisect qualified+import Tests.Extra.DsuMonoid qualified import Tests.Extra.DynLazySegTree qualified import Tests.Extra.DynLazySegTree.Persistent qualified import Tests.Extra.DynSegTree qualified@@ -67,7 +69,9 @@         testGroup "Dsu" Tests.Dsu.tests,         testGroup           "Extra"-          [ testGroup "Bisect" Tests.Extra.Bisect.tests,+          [ testGroup "AhoCorasick" Tests.Extra.AhoCorasick.tests,+            testGroup "Bisect" Tests.Extra.Bisect.tests,+            testGroup "DsuMonoid" Tests.Extra.DsuMonoid.tests,             testGroup "DynLazySegTree" Tests.Extra.DynLazySegTree.tests,             testGroup "DynLazySegTree.Persistent" Tests.Extra.DynLazySegTree.Persistent.tests,             testGroup "DynSegTree" Tests.Extra.DynSegTree.tests,

@@ -0,0 +1,171 @@+module Tests.Extra.Vector where++import AtCoder.Extra.Vector qualified as EV+import Control.Monad.ST (runST)+import Data.Functor.Identity (Identity, runIdentity)+import Data.List qualified as L+import Data.Vector qualified as V+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU+import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck as QC++prop_argsort :: [Int] -> QC.Property+prop_argsort xs =+  let lhs = VU.fromList . map snd . L.sort $ zip xs [0 :: Int ..]+      rhs = EV.argsort $ VU.fromList xs+   in lhs QC.=== rhs++prop_concatMapM :: [Int] -> QC.Property+prop_concatMapM xs =+  let f x = VU.fromList [x, x, x]+      vec = VU.fromList xs+      lhs = VU.concatMap f vec+      rhs = runST $ EV.concatMapM (pure . f) vec+   in lhs QC.=== rhs++prop_iconcatMap :: [Int] -> QC.Property+prop_iconcatMap xs =+  let f i x = VU.fromList [i + x, i + x, i + x]+      vec = VU.fromList xs+      lhs = VU.concat $ zipWith f [0 :: Int ..] xs+      rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec+   in lhs QC.=== rhs++prop_iconcatMapM :: [Int] -> QC.Property+prop_iconcatMapM xs =+  let f i x = VU.fromList [i + x, i + x, i + x]+      vec = VU.fromList xs+      lhs = VU.concat $ zipWith f [0 :: Int ..] xs+      rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec+   in lhs QC.=== rhs++prop_mapAccumL :: [Int] -> QC.Property+prop_mapAccumL xs =+  let f s x = (s * x, s + x)+      (!l1, !l2) = L.mapAccumL f (0 :: Int) xs+      (!r1, !r2) = EV.mapAccumL f (0 :: Int) $ VU.fromList xs+   in QC.conjoin [l1 QC.=== r1, VU.fromList l2 QC.=== r2]++-- | scanM etc.+prop_monadicScanlLike ::+  ((Int -> Int -> Int) -> Int -> VU.Vector Int -> VU.Vector Int) ->+  ((Int -> Int -> Identity Int) -> Int -> VU.Vector Int -> Identity (VU.Vector Int)) ->+  Int ->+  [Int] ->+  QC.Property+prop_monadicScanlLike ref acl x xs =+  let xs' = VU.fromList xs+      f = (+)+      mf x y = pure $ x + y+   in ref f x xs' QC.=== runIdentity (acl mf x xs')++-- | scanM1 etc.+prop_monadicScanl1Like ::+  ((Int -> Int -> Int) -> VU.Vector Int -> VU.Vector Int) ->+  ((Int -> Int -> Identity Int) -> VU.Vector Int -> Identity (VU.Vector Int)) ->+  QC.NonEmptyList Int ->+  QC.Property+prop_monadicScanl1Like ref acl (QC.NonEmpty xs) =+  let xs' = VU.fromList xs+      f = (+)+      mf x y = pure $ x + y+   in ref f xs' QC.=== runIdentity (acl mf xs')++prop_chunks :: QC.Positive Int -> [Int] -> QC.Property+prop_chunks (QC.Positive k) [] = EV.chunks k (VU.empty @Int) QC.=== V.empty+prop_chunks (QC.Positive k) xs =+  let res = EV.chunks k $ VU.fromList xs+      n = length xs+   in QC.conjoin+        [ V.sum (VG.map VG.length res) QC.=== n,+          V.all ((== k) . VG.length) (V.init res) QC.=== True,+          VG.concat (V.toList res) QC.=== VU.fromList xs+        ]++unit_maxRangeSum :: TestTree+unit_maxRangeSum = testCase "unit_maxRangeSum" $ do+  EV.maxRangeSum (VU.singleton (-1 :: Int)) @?= 0+  EV.maxRangeSum (VU.empty @Int) @?= 0++prop_maxRangeSum :: [Int] -> QC.Property+prop_maxRangeSum xs =+  let vec = VU.fromList xs+      lhs =+        let n = VU.length vec+            lrs = [(l, r) | l <- [0 .. n], r <- [l .. n]]+            eval (!l, !r) = VU.sum . VU.take (r - l) $ VU.drop l vec+         in maximum $ map eval lrs+      rhs = EV.maxRangeSum vec+   in lhs QC.=== rhs++unit_minRangeSum :: TestTree+unit_minRangeSum = testCase "unit_minRangeSum [1]" $ do+  EV.minRangeSum (VU.singleton (1 :: Int)) @?= 0+  EV.minRangeSum (VU.empty @Int) @?= 0++prop_minRangeSum :: [Int] -> QC.Property+prop_minRangeSum xs =+  let vec = VU.fromList xs+      lhs =+        let n = VU.length vec+            lrs = [(l, r) | l <- [0 .. n], r <- [l .. n]]+            eval (!l, !r) = VU.sum . VU.take (r - l) $ VU.drop l vec+         in minimum $ map eval lrs+      rhs = EV.minRangeSum vec+   in lhs QC.=== rhs++prop_slideMinIndices :: QC.Positive Int -> [Int] -> QC.Property+prop_slideMinIndices (QC.Positive k) xs =+  let vec = VU.fromList xs+   in slideMin k vec QC.=== EV.slideMinIndices k vec++prop_slideMaxIndices :: QC.Positive Int -> [Int] -> QC.Property+prop_slideMaxIndices (QC.Positive k) xs =+  let vec = VU.fromList xs+   in slideMax k vec QC.=== EV.slideMaxIndices k vec++slideMin :: Int -> VU.Vector Int -> VU.Vector Int+slideMin k xs+  | VU.null xs = VU.empty+  | k >= n = VU.singleton $ VU.minIndex xs+  | otherwise = VU.generate (n - (k - 1)) $ \l ->+      let slice = VU.take k $ VU.drop l xs+       in (+ l) $ VU.minIndex slice+  where+    n = VU.length xs++slideMax :: Int -> VU.Vector Int -> VU.Vector Int+slideMax k xs+  | VU.null xs = VU.empty+  | k >= n = VU.singleton $ VU.maxIndex xs+  | otherwise = VU.generate (n - (k - 1)) $ \l ->+      let slice = VU.take k $ VU.drop l xs+       in (+ l) $ VU.maxIndex slice+  where+    n = VU.length xs++tests :: [TestTree]+tests =+  [ QC.testProperty "argsort" prop_argsort,+    QC.testProperty "concatMapM" prop_concatMapM,+    QC.testProperty "iconcatMap" prop_iconcatMap,+    QC.testProperty "iconcatMapM" prop_iconcatMapM,+    QC.testProperty "mapAccumL" prop_mapAccumL,+    QC.testProperty "chunks" prop_chunks,+    QC.testProperty "prescanlM" (prop_monadicScanlLike VU.prescanl EV.prescanlM),+    QC.testProperty "prescanlM'" (prop_monadicScanlLike VU.prescanl' EV.prescanlM'),+    QC.testProperty "postscanlM" (prop_monadicScanlLike VU.postscanl EV.postscanlM),+    QC.testProperty "postscanlM'" (prop_monadicScanlLike VU.postscanl' EV.postscanlM'),+    QC.testProperty "scanlM" (prop_monadicScanlLike VU.scanl EV.scanlM),+    QC.testProperty "scanlM'" (prop_monadicScanlLike VU.scanl' EV.scanlM'),+    QC.testProperty "scanl1M" (prop_monadicScanl1Like VU.scanl1 EV.scanl1M),+    QC.testProperty "scanl1M'" (prop_monadicScanl1Like VU.scanl1' EV.scanl1M'),+    QC.testProperty "maxRangeSum" prop_maxRangeSum,+    QC.testProperty "minRangeSum" prop_minRangeSum,+    unit_maxRangeSum,+    unit_minRangeSum,+    QC.testProperty "slideMinIndices" prop_slideMinIndices,+    QC.testProperty "slideMaxIndices" prop_slideMaxIndices+  ]

@@ -0,0 +1,152 @@+{-# LANGUAGE ViewPatterns #-}++module Tests.Extra.Vector.Prim where++import AtCoder.Extra.Vector.Prim qualified as EV+import Control.Monad.ST (ST, runST)+import Data.Vector.Generic qualified as VG+import Data.Vector.Unboxed qualified as VU+import Test.Tasty+import Test.Tasty.QuickCheck as QC++prop_replicateM :: Int -> Int -> QC.Property+prop_replicateM n x = VU.replicate n x QC.=== runST (EV.replicateM n (pure x))++prop_generateM :: Int -> Int -> QC.Property+prop_generateM n x = VU.generate n (+ x) QC.=== runST (EV.generateM n (pure . (+ x)))++prop_iterateNM :: Int -> Int -> QC.Property+prop_iterateNM n x = VU.iterateN n (* 2) x QC.=== runST (EV.iterateNM n (pure . (* 2)) x)++prop_constructNM :: QC.NonNegative Int -> QC.Property+prop_constructNM (QC.NonNegative n) = VU.constructN n VG.length QC.=== runST (EV.constructNM n (pure . VG.length))++prop_constructrNM :: QC.NonNegative Int -> QC.Property+prop_constructrNM (QC.NonNegative n) = VU.constructrN n VG.length QC.=== runST (EV.constructrNM n (pure . VG.length))++prop_mapM :: [Int] -> QC.Property+prop_mapM (VU.fromList -> xs) = VU.map (* 2) xs QC.=== runST (EV.mapM (pure . (* 2)) xs)++-- prop_mapM_ :: [Int] -> QC.Property+-- prop_mapM_ xs =++prop_imapM :: [Int] -> QC.Property+prop_imapM (VU.fromList -> xs) = VU.imap (\i x -> 2 * (i + x)) xs QC.=== runST (EV.imapM (\i x -> pure (2 * (i + x))) xs)++-- prop_imapM_ :: [Int] -> QC.Property+-- prop_imapM_ xs =++prop_iforM :: [Int] -> QC.Property+prop_iforM (VU.fromList -> xs) = VU.imap (\i x -> 2 * (i + x)) xs QC.=== runST (EV.iforM xs (\i x -> pure (2 * (i + x))))++-- prop_iforM_ :: [Int] -> QC.Property+-- prop_iforM_ xs =++prop_zipWithM :: [Int] -> [Int] -> QC.Property+prop_zipWithM (VU.fromList -> xs) (VU.fromList -> ys) =+  VU.zipWith (*) xs ys QC.=== runST (EV.zipWithM (\x y -> pure (x * y)) xs ys)++-- prop_zipWithM_ :: [Int] -> [Int] -> QC.Property+-- prop_zipWithM_  (VU.fromList -> xs) (VU.fromList -> ys) =++prop_izipWithM :: [Int] -> [Int] -> QC.Property+prop_izipWithM (VU.fromList -> xs) (VU.fromList -> ys) =+  VU.izipWith (\i a b -> i + a * b) xs ys QC.=== runST (EV.izipWithM (\i a b -> pure (i + a * b)) xs ys)++-- prop_izipWithM_ :: [Int] -> [Int] -> QC.Property+-- prop_izipWithM_  (VU.fromList -> xs) (VU.fromList -> ys) =++prop_concatMapM :: [Int] -> QC.Property+prop_concatMapM xs =+  let f x = VU.fromList [x, x, x]+      vec = VU.fromList xs+      lhs = VU.concatMap f vec+      rhs = runST $ EV.concatMapM (pure . f) vec+   in lhs QC.=== rhs++prop_iconcatMap :: [Int] -> QC.Property+prop_iconcatMap xs =+  let f i x = VU.fromList [i + x, i + x, i + x]+      vec = VU.fromList xs+      lhs = VU.concat $ zipWith f [0 :: Int ..] xs+      rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec+   in lhs QC.=== rhs++prop_iconcatMapM :: [Int] -> QC.Property+prop_iconcatMapM xs =+  let f i x = VU.fromList [i + x, i + x, i + x]+      vec = VU.fromList xs+      lhs = VU.concat $ zipWith f [0 :: Int ..] xs+      rhs = runST $ EV.iconcatMapM (\i x -> pure (f i x)) vec+   in lhs QC.=== rhs++prop_filterM :: [Int] -> QC.Property+prop_filterM (VU.fromList -> xs) = VU.filter odd xs QC.=== runST (EV.filterM (pure . odd) xs)++prop_mapMaybeM :: [Int] -> QC.Property+prop_mapMaybeM (VU.fromList -> xs) =+  let f x = if odd x then Just (2 * x) else Nothing+   in VU.mapMaybe f xs QC.=== runST (EV.mapMaybeM (pure . f) xs)++prop_imapMaybeM :: [Int] -> QC.Property+prop_imapMaybeM (VU.fromList -> xs) =+  let f i x = if odd i then Just (i * x) else Nothing+   in VU.imapMaybe f xs QC.=== runST (EV.imapMaybeM (\i -> pure . f i) xs)++-- | scanM etc.+prop_monadicScanlLike ::+  ((Int -> Int -> Int) -> Int -> VU.Vector Int -> VU.Vector Int) ->+  (forall s. (Int -> Int -> ST s Int) -> Int -> VU.Vector Int -> ST s (VU.Vector Int)) ->+  Int ->+  [Int] ->+  QC.Property+prop_monadicScanlLike ref acl x xs =+  let xs' = VU.fromList xs+      f = (+)+      mf x y = pure $ x + y+   in ref f x xs' QC.=== runST (acl mf x xs')++-- | scanM1 etc.+prop_monadicScanl1Like ::+  ((Int -> Int -> Int) -> VU.Vector Int -> VU.Vector Int) ->+  (forall s. (Int -> Int -> ST s Int) -> VU.Vector Int -> ST s (VU.Vector Int)) ->+  QC.NonEmptyList Int ->+  QC.Property+prop_monadicScanl1Like ref acl (QC.NonEmpty xs) =+  let xs' = VU.fromList xs+      f = (+)+      mf x y = pure $ x + y+   in ref f xs' QC.=== runST (acl mf xs')++tests :: [TestTree]+tests =+  [ QC.testProperty "replicateM" prop_replicateM,+    QC.testProperty "generateM" prop_generateM,+    QC.testProperty "iterateNM" prop_iterateNM,+    QC.testProperty "constructNM" prop_constructNM,+    QC.testProperty "constructrNM" prop_constructrNM,+    QC.testProperty "mapM" prop_mapM,+    -- QC.testProperty "mapM_" prop_mapM_,+    QC.testProperty "imapM" prop_imapM,+    -- QC.testProperty "imapM_" prop_imapM_,+    QC.testProperty "iforM" prop_iforM,+    -- QC.testProperty "iforM_" prop_iforM_,+    QC.testProperty "zipWithM" prop_zipWithM,+    -- QC.testProperty "zipWithM_" prop_zipWithM_,+    QC.testProperty "izipWithM" prop_izipWithM,+    -- QC.testProperty "izipWithM_" prop_izipWithM_,+    QC.testProperty "concatMapM" prop_concatMapM,+    QC.testProperty "iconcatMap" prop_iconcatMap,+    QC.testProperty "iconcatMapM" prop_iconcatMapM,+    QC.testProperty "filterM" prop_filterM,+    QC.testProperty "mapMaybeM" prop_mapMaybeM,+    QC.testProperty "imapMaybeM" prop_imapMaybeM,+    QC.testProperty "prescanlM" (prop_monadicScanlLike VU.prescanl EV.prescanlM),+    QC.testProperty "prescanlM'" (prop_monadicScanlLike VU.prescanl' EV.prescanlM'),+    QC.testProperty "postscanlM" (prop_monadicScanlLike VU.postscanl EV.postscanlM),+    QC.testProperty "postscanlM'" (prop_monadicScanlLike VU.postscanl' EV.postscanlM'),+    QC.testProperty "scanlM" (prop_monadicScanlLike VU.scanl EV.scanlM),+    QC.testProperty "scanlM'" (prop_monadicScanlLike VU.scanl' EV.scanlM'),+    QC.testProperty "scanl1M" (prop_monadicScanl1Like VU.scanl1 EV.scanl1M),+    QC.testProperty "scanl1M'" (prop_monadicScanl1Like VU.scanl1' EV.scanl1M')+  ]