packages feed

ac-library-hs 1.3.1.0 → 1.4.0.0

raw patch · 26 files changed

+637/−65 lines, 26 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

+ AtCoder.Extra.Graph: findCycleDirected :: (HasCallStack, Unbox w) => Csr w -> Maybe (Vector Int, Vector Int)
+ AtCoder.Extra.Graph: findCycleUndirected :: (HasCallStack, Unbox w) => Csr w -> Maybe (Vector Int, Vector Int)
+ AtCoder.Extra.Pdsu: mergeMaybe :: (HasCallStack, PrimMonad m, Monoid a, Ord a, Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m (Maybe Int)
+ AtCoder.Internal.Csr: eAdjW :: (HasCallStack, Unbox w) => Csr w -> Int -> Vector (Int, Int, w)
+ AtCoder.Internal.GrowVec: clear :: (PrimMonad m, Unbox a) => GrowVec (PrimState m) a -> m ()
+ AtCoder.Internal.GrowVec: readBack :: (HasCallStack, PrimMonad m, Unbox a) => GrowVec (PrimState m) a -> Int -> m a
+ AtCoder.Internal.GrowVec: readBackMaybe :: (HasCallStack, PrimMonad m, Unbox a) => GrowVec (PrimState m) a -> Int -> m (Maybe a)
+ AtCoder.Internal.GrowVec: reverse :: (HasCallStack, PrimMonad m, Ord a, Unbox a) => GrowVec (PrimState m) a -> m ()
- AtCoder.Extra.Pdsu: merge :: (HasCallStack, PrimMonad m, Monoid a, Ord a, Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool
+ AtCoder.Extra.Pdsu: merge :: (HasCallStack, PrimMonad m, Monoid a, Ord a, Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Int

Files

CHANGELOG.md view
@@ -1,5 +1,16 @@ # Revision history for acl-hs +## 1.4.0.0 -- May 2025++- Added cycle detection functions to `AtCoder.Extra.Graph`+- Fixed the exteremely slow `Internal.GrowVec` functions with `{-# INLINE #-}`+- Changed `AtCoder.Extra.Pdsu.merge` to return new representative.+- Fixed `AtCoder.Extra.Hld.jump` to return `Nothing` for `k < 0`.++## 1.3.1.0 -- May 2025++- Fixed `AtCoder.Extra.Monoid.RollingHash` to take 64 bit mod values.+ ## 1.3.0.0 -- April 2025  - Added `AtCoder.Extra.Math.isPrimitiveRoot`.
ac-library-hs.cabal view
@@ -4,7 +4,7 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.3.1.0+version:         1.4.0.0 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive@@ -162,6 +162,7 @@     Tests.Extra.ModInt64     Tests.Extra.Monoid     Tests.Extra.MultiSet+    Tests.Extra.Pdsu     Tests.Extra.SegTree2d     Tests.Extra.SegTree2d.Dense     Tests.Extra.Semigroup.Matrix@@ -224,6 +225,7 @@     Bench.AddMod     Bench.Matrix     Bench.ModInt+    Bench.Montgomery64     Bench.PowMod     Bench.RepeatWithIndex     Bench.RepeatWithoutIndex@@ -235,6 +237,8 @@     BenchLib.MulMod.Barrett64     BenchLib.MulMod.BarrettWideWord     BenchLib.MulMod.Montgomery+    BenchLib.Montgomery64.Inline+    BenchLib.Montgomery64.Noinline     BenchLib.PowMod     BenchLib.SwapDupe 
+ benchmarks/Bench/Montgomery64.hs view
@@ -0,0 +1,41 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE MagicHash #-}++module Bench.Montgomery64 (benches) where++import BenchLib.Montgomery64.Inline qualified as M64A+import BenchLib.Montgomery64.Noinline qualified as M64B+import Criterion+import Data.Vector.Unboxed qualified as VU+import GHC.Exts (proxy#)+import System.Random++benches :: Benchmark+benches =+  bgroup+    "Montgomery64"+    [ -- The following four are actually meaningless:+      bench "new @998244353 inline" $ whnf (VU.foldl' (\ !_ !_ -> M64A.new (proxy# @998244353)) mA) randomVec,+      bench "new @998244353 noinline" $ whnf (VU.foldl' (\ !_ !_ -> M64B.new (proxy# @998244353)) mB) randomVec,+      bench "new @2305843009213693951 inline" $ whnf (VU.foldl' (\ !_ !_ -> M64A.new (proxy# @2305843009213693951)) mA) randomVec,+      bench "new @2305843009213693951 noinline" $ whnf (VU.foldl' (\ !_ !_ -> M64B.new (proxy# @2305843009213693951)) mB) randomVec,+      -- These are the important tests:+      bench "fromVal 998244353 inline" $ whnf (VU.foldl' (\ !_ !x -> M64A.fromVal (fromIntegral x)) mA) repVec1,+      bench "fromVal 998244353 noinline" $ whnf (VU.foldl' (\ !_ !x -> M64B.fromVal (fromIntegral x)) mB) repVec1,+      bench "fromVal 2305843009213693951 inline" $ whnf (VU.foldl' (\ !_ !x -> M64A.fromVal (fromIntegral x)) mA) repVec2,+      bench "fromVal 2305843009213693951 noinline" $ whnf (VU.foldl' (\ !_ !x -> M64B.fromVal (fromIntegral x)) mB) repVec2,+      bench "fromVal random inline" $ whnf (VU.foldl' (\ !_ !x -> M64A.fromVal (fromIntegral x)) mA) randomVec,+      bench "fromVal random noinline" $ whnf (VU.foldl' (\ !_ !x -> M64B.fromVal (fromIntegral x)) mB) randomVec+    ]+  where+    n = 10000+    mA = M64A.fromVal 3+    mB = M64B.fromVal 3+    randomVec :: VU.Vector Int+    randomVec =+      VU.map fromIntegral $+        VU.unfoldrExactN n ((\(!x, !gen) -> (1 + 2 * x, gen)) <$> genWord64R (2 ^ 62 `div` 2)) (mkStdGen 123456789)+    repVec1 :: VU.Vector Int+    repVec1 = VU.replicate n 998244353+    repVec2 :: VU.Vector Int+    repVec2 = VU.replicate n 2305843009213693951
+ benchmarks/BenchLib/Montgomery64/Inline.hs view
@@ -0,0 +1,40 @@+{-# LANGUAGE MagicHash #-}++module BenchLib.Montgomery64.Inline (Montgomery64 (..), new, fromVal) where++import AtCoder.Internal.Assert qualified as ACIA+import Data.Bits (bit, (!>>.))+import Data.WideWord.Word128 (Word128 (..))+import Data.Word (Word64)+import GHC.Exts (Proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal')++data Montgomery64 = Montgomery64+  { mM64 :: {-# UNPACK #-} !Word64,+    rM64 :: {-# UNPACK #-} !Word64,+    n2M64 :: {-# UNPACK #-} !Word64+  }+  deriving+    ( Eq,+      Show+    )++{-# INLINE new #-}+new :: forall a. (HasCallStack, KnownNat a) => Proxy# a -> Montgomery64+new p = fromVal . fromIntegral $! natVal' p++{-# INLINE fromVal #-}+fromVal :: (HasCallStack) => Word64 -> Montgomery64+fromVal m =+  let !m128 :: Word128 = fromIntegral m+      !n2 = word128Lo64 $ (-m128) `mod` m128+      !r = getR m 0+      !_ = ACIA.runtimeAssert (r * m == -1) "AtCoder.Extra.Montgomery64.fromVal: internal implementation error"+   in Montgomery64 m r n2+  where+    !_ = ACIA.runtimeAssert (odd m && m <= bit 62) $ "AtCoder.Extra.Montgomery64.fromVal: not given odd modulus value that is less than or equal to 2^62: " ++ show m+    getR :: Word64 -> Int -> Word64+    getR !acc i+      | i >= 5 = -acc+      | otherwise = getR (acc * (2 - m * acc)) (i + 1)
+ benchmarks/BenchLib/Montgomery64/Noinline.hs view
@@ -0,0 +1,40 @@+{-# LANGUAGE MagicHash #-}++module BenchLib.Montgomery64.Noinline (Montgomery64 (..), new, fromVal) where++import AtCoder.Internal.Assert qualified as ACIA+import Data.Bits (bit, (!>>.))+import Data.WideWord.Word128 (Word128 (..))+import Data.Word (Word64)+import GHC.Exts (Proxy#)+import GHC.Stack (HasCallStack)+import GHC.TypeNats (KnownNat, natVal')++data Montgomery64 = Montgomery64+  { mM64 :: {-# UNPACK #-} !Word64,+    rM64 :: {-# UNPACK #-} !Word64,+    n2M64 :: {-# UNPACK #-} !Word64+  }+  deriving+    ( Eq,+      Show+    )++{-# INLINE new #-}+new :: forall a. (HasCallStack, KnownNat a) => Proxy# a -> Montgomery64+new p = fromVal . fromIntegral $! natVal' p++{-# NOINLINE fromVal #-}+fromVal :: (HasCallStack) => Word64 -> Montgomery64+fromVal m =+  let !m128 :: Word128 = fromIntegral m+      !n2 = word128Lo64 $ (-m128) `mod` m128+      !r = getR m 0+      !_ = ACIA.runtimeAssert (r * m == -1) "AtCoder.Extra.Montgomery64.fromVal: internal implementation error"+   in Montgomery64 m r n2+  where+    !_ = ACIA.runtimeAssert (odd m && m <= bit 62) $ "AtCoder.Extra.Montgomery64.fromVal: not given odd modulus value that is less than or equal to 2^62: " ++ show m+    getR :: Word64 -> Int -> Word64+    getR !acc i+      | i >= 5 = -acc+      | otherwise = getR (acc * (2 - m * acc)) (i + 1)
benchmarks/Main.hs view
@@ -5,6 +5,7 @@ import Bench.AddMod qualified import Bench.Matrix qualified import Bench.ModInt qualified+import Bench.Montgomery64 qualified import Bench.MulMod qualified import Bench.PowMod qualified import Bench.RepeatWithIndex qualified@@ -12,9 +13,6 @@ import Bench.SwapDupe qualified import Criterion.Main --- TODO: try tasty-bench-#define MOD 998244353- main :: IO () main =   defaultMain@@ -24,6 +22,7 @@       Bench.AddMod.benches,       Bench.PowMod.benches,       Bench.Matrix.benches,+      Bench.Montgomery64.benches,       Bench.RepeatWithIndex.benches,       Bench.RepeatWithoutIndex.benches,       Bench.SwapDupe.benches
src/AtCoder/Dsu.hs view
@@ -23,6 +23,9 @@ -- >>> Dsu.merge dsu 0 1  -- 0=1 2 3 -- 0 --+-- >>> Dsu.mergeMaybe dsu 0 1 -- already merged+-- Nothing+-- -- >>> Dsu.merge_ dsu 1 2 -- 0=1=2 3 -- >>> Dsu.mergeMaybe dsu 1 2 -- Nothing
src/AtCoder/Extra/Bisect.hs view
@@ -41,11 +41,11 @@ import Data.Vector.Generic qualified as VG import GHC.Stack (HasCallStack) --- | \(O(\log n)\) Returns the maximum \(r\) where \(x_i \lt x_0\) holds for \(i \in [0, r)\).+-- | \(O(\log n)\) Returns the maximum \(r\) where \(x_i \lt x_{ref}\) holds for \(i \in [0, r)\). -- -- @--- Y Y Y Y Y N N N N N      Y: x_i < x_0--- --------- *---------> x  N: x_i >= x_0+-- Y Y Y Y Y N N N N N      Y: x_i < x_ref+-- --------- *---------> x  N: x_i >= x_ref --           R              R: the right boundary point returned -- @ --@@ -104,11 +104,11 @@   where     !_ = ACIA.checkIntervalBounded "AtCoder.Extra.Bisect.lowerBoundIn" l r $ VG.length vec --- | \(O(\log n)\) Returns the maximum \(r\) where \(x_i \le x_0\) holds for \(i \in [0, r)\).+-- | \(O(\log n)\) Returns the maximum \(r\) where \(x_i \le x_{ref}\) holds for \(i \in [0, r)\). -- -- @--- Y Y Y Y Y N N N N N      Y: x_i <= x_0,--- --------- *---------> x  N: x_i > x_0,+-- Y Y Y Y Y N N N N N      Y: x_i <= x_ref,+-- --------- *---------> x  N: x_i > x_ref, --           R              R: the right boundary point returned -- @ --
src/AtCoder/Extra/DynLazySegTree.hs view
@@ -62,7 +62,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Applications     applyAt,
src/AtCoder/Extra/DynLazySegTree/Persistent.hs view
@@ -65,7 +65,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Applications     applyAt,
src/AtCoder/Extra/DynSegTree.hs view
@@ -56,7 +56,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Tree operations     resetInterval,
src/AtCoder/Extra/DynSegTree/Persistent.hs view
@@ -54,7 +54,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Tree operations     resetInterval,
src/AtCoder/Extra/DynSparseSegTree.hs view
@@ -55,7 +55,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Binary searches     maxRight,
src/AtCoder/Extra/DynSparseSegTree/Persistent.hs view
@@ -54,7 +54,7 @@     -- * Products     prod,     -- prodMaybe,-    allProd, -- FIXME: rename it to prodAll+    allProd,      -- * Binary searches     maxRight,
src/AtCoder/Extra/Graph.hs view
@@ -8,7 +8,7 @@ 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.+    -- | The @Csr@ data type and all the functions such as `build` or `adj` are re-exported.     -- See the @Csr@ module for details.     module Csr, @@ -17,6 +17,8 @@     swapDupe',     scc,     rev,+    findCycleDirected,+    findCycleUndirected,      -- * Generic graph functions @@ -99,25 +101,31 @@ where  import AtCoder.Dsu qualified as Dsu+import AtCoder.Extra.HashMap qualified as HM import AtCoder.Extra.IntSet qualified as IS import AtCoder.Extra.Ix0 (Bounds0, Ix0 (..))+import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Buffer qualified as B import AtCoder.Internal.Csr as Csr+import AtCoder.Internal.GrowVec qualified as GV import AtCoder.Internal.MinHeap qualified as MH import AtCoder.Internal.Queue qualified as Q import AtCoder.Internal.Scc qualified as ACISCC-import Control.Monad (replicateM_, when)+import Control.Applicative ((<|>))+import Control.Monad (replicateM_, unless, when) import Control.Monad.Fix (fix) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Control.Monad.ST (ST, runST) import Data.Bit (Bit (..))+import Data.Bits ((.<<.), (.|.)) import Data.Foldable (for_)-import Data.Maybe (fromJust)+import Data.Maybe (fromJust, fromMaybe) 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 Data.Word (Word8) import GHC.Stack (HasCallStack)  -- | \(O(n)\) Converts directed edges into non-directed edges; each edge \((u, v, w)\) is duplicated@@ -231,6 +239,281 @@           !o2 = startCsr VG.! (v1 + 1)           !vw2s = VU.slice o1 (o2 - o1) vws        in VU.map (\(!v2, !w2) -> (v2, v1, w2)) vw2s++-- TODO: is this minimum cycle?++-- | \(O(n + m)\) Given a directed graph, finds a minimal cycle and returns a vector of vertices and+-- a vector of @(vertices, csrEdgeIndices)@.+--+-- ==== __Example__+--+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let gr = Gr.build' 4 $ VU.fromList [(0, 1), (1, 2), (2, 3), (3, 1)]+-- >>> findCycleDirected gr -- returns (vs, es)+-- Just ([1,2,3],[1,2,3])+--+-- @since 1.4.0.0+{-# INLINEABLE findCycleDirected #-}+findCycleDirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)+findCycleDirected gr@Csr {..} = runST $ do+  used <- VUM.replicate @_ @Word8 nCsr 0+  -- par <- VUM.unsafeNew @_ @(Int, Int) nCsr+  par <- VUM.replicate nCsr (-1 :: Int, -1 :: Int)+  vs <- GV.new @_ @Int 16+  es <- GV.new @_ @Int 16+  esFrom <- GV.new @_ @Int 16 -- If we had `from` in Csr, we could skip this+  let dfs u = do+        VGM.write used u 1+        let next evs = case VU.uncons evs of+              Nothing -> pure ()+              Just ((!iEdge, !v), !evs') -> do+                b <- GV.null es+                when b $ do+                  use <- VGM.read used v+                  case use of+                    0 -> do+                      VGM.write par v (u, iEdge)+                      dfs v+                      next evs'+                    1 -> do+                      GV.pushBack es iEdge+                      GV.pushBack esFrom u+                      let backtrack cur+                            | cur == v = pure ()+                            | otherwise = do+                                (!prevVert, !edge) <- VGM.read par cur+                                GV.pushBack es edge+                                GV.pushBack esFrom prevVert+                                backtrack prevVert+                      backtrack u+                      GV.reverse es+                      GV.reverse esFrom+                    _ -> do+                      next evs'++        next $ eAdj gr u+        VGM.write used u 2++  VGM.iforM_ used $ \v use -> do+    when (use == 0) $ do+      dfs v++  b <- GV.null es+  unless b $ do+    -- find minimum cycle+    nxt <- VUM.replicate nCsr (-1 :: Int) -- edge indices+    do+      es' <- GV.unsafeFreeze es+      esFrom' <- GV.unsafeFreeze esFrom+      VU.forM_ (VU.zip es' esFrom') $ \(!iEdge, !vFrom) -> do+        VGM.write nxt vFrom iEdge++    for_ [0 .. nCsr - 1] $ \vA -> do+      nxtA <- VGM.read nxt vA+      unless (nxtA == -1) $ do+        VU.forM_ (eAdj gr vA) $ \(!iEdge, !vB) -> do+          nxtB <- VGM.read nxt vB+          unless (nxtB == -1 || adjCsr VG.! nxtA == vB) $ do+            let inner x+                  | x == vB = pure ()+                  | otherwise = do+                      nxtX <- VGM.exchange nxt x (-1)+                      inner $ adjCsr VG.! nxtX+            inner vA+            VGM.write nxt vA iEdge++    GV.clear es+    let loop v+          | v >= nCsr = pure ()+          | otherwise = do+              nxtV <- VGM.read nxt v+              if nxtV == -1+                then loop (v + 1)+                else do+                  let inner x = do+                        GV.pushBack vs x+                        nxtX <- VGM.read nxt x+                        GV.pushBack es nxtX+                        let !x' = adjCsr VG.! nxtX+                        unless (x' == v) $ inner x'+                  inner v+    loop 0++  vs' <- GV.unsafeFreeze vs+  es' <- GV.unsafeFreeze es+  if VU.null es'+    then pure Nothing+    else pure $ Just (vs', es')++-- | \(O(n + m)\) Given an undirected graph, finds a minimal cycle and returns a vector of vertices+-- a vector of @(vertices, csrEdgeIndices)@. A single edge index does not make much sense for an+-- undirected graph, so map back to the original edge index manually if needed.+--+-- ==== Constraints+-- - The graph must be created with `swapDupe` or `swapDupe'`. Otherwise the returned edge indices+--   could make no sense.+--+-- ==== __Example__+--+-- >>> import AtCoder.Extra.Graph qualified as Gr+-- >>> import Data.Vector.Unboxed qualified as VU+-- >>> let gr = Gr.build' 4 . Gr.swapDupe' $ VU.fromList [(0, 1), (1, 2), (1, 3), (2, 3)]+-- >>> findCycleUndirected gr -- returns (vs, es)+-- Just ([1,3,2],[3,5,2])+--+-- Retrieve original edge indices that makes up the cycle, by recording them in edge weights:+--+-- >>> let gr = Gr.build 4 . Gr.swapDupe $ VU.fromList [(0, 1, 0 :: Int), (1, 2, 1), (1, 3, 2), (2, 3, 3)]+-- >>> let Just (vs, es) = findCycleUndirected gr -- returns (vs, es)+-- >>> VU.backpermute (Gr.wCsr gr) es+-- [2,3,1]+--+-- It's a bit hacky.+--+-- @since 1.4.0.0+{-# INLINEABLE findCycleUndirected #-}+findCycleUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)+findCycleUndirected gr@Csr {..} =+  let !_ = ACIA.runtimeAssert (even mCsr) $ "AtCoder.Extra.Graph.findCycleUndirected: the number of edge in an undirected graph must be even: `" ++ show mCsr ++ "`"+   in -- If we have the same edge id for duplicated edges, `findCycleSimpleUndirected` could be modified+      -- to handle both complex and simple graph. We don't, and we need the complex graph handling.+      -- This is not optimal, but we need a dedicated `buildUndirected` function and different edge ID+      -- (not index) handling in CSR if we go with the optimal approach.+      --+      --  Note that the implementations are suspecious..+      findCycleComplexUndirected gr <|> findCycleSimpleUndirected gr++{-# INLINEABLE findCycleComplexUndirected #-}+findCycleComplexUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)+findCycleComplexUndirected gr@Csr {..} = runST $ do+  usedE <- HM.new @_ @Int (mCsr `div` 2 + {- not needed, but in case of panic? -} 4)+  cntE <- HM.new @_ @Word8 (mCsr `div` 2 + {- not needed, but in case of panic? -} 4)++  -- we'll give unique indices to (u, v) pairs+  let ix u v = min u v .<<. 32 .|. max u v++  let nextU u+        | u >= nCsr = pure Nothing+        | otherwise = do+            let nextV evs = case VU.uncons evs of+                  Nothing -> pure Nothing+                  Just ((!e, !v), !evs') -> case compare u v of+                    -- self loop edge+                    EQ -> pure $ Just (VU.singleton v, VU.singleton e)+                    LT -> do+                      let !i = ix u v+                      c <- fromMaybe 0 <$> HM.lookup cntE i+                      case c of+                        0 -> do+                          HM.insert usedE i e+                          HM.insert cntE i 1+                          nextV evs'+                        1 -> do+                          -- found the first duplicated edge+                          HM.insert cntE i 2+                          nextV evs'+                        _ -> do+                          nextV evs'+                    GT -> do+                      let !i = ix u v+                      cnt <- fromMaybe 0 <$> HM.lookup cntE i+                      case cnt of+                        2 -> do+                          -- there are duplicate edges between (u, v) and this is the+                          -- first (u, v) (u > v)+                          HM.insert cntE i 3+                          nextV evs'+                        3 -> do+                          -- this is the second duplicate edge (u, v) (u > v)+                          e1 <- fromJust <$> HM.lookup usedE i+                          let vs = VU.fromListN 2 [v, u]+                          let es = VU.fromListN 2 [e1, e]+                          pure $ Just (vs, es)+                        _ -> nextV evs'++            res <- nextV $ eAdj gr u+            case res of+              Just ret -> pure $ Just ret+              Nothing -> nextU (u + 1)++  nextU 0++{-# INLINEABLE findCycleSimpleUndirected #-}+findCycleSimpleUndirected :: (HasCallStack, VU.Unbox w) => Csr w -> Maybe (VU.Vector Int, VU.Vector Int)+findCycleSimpleUndirected gr@Csr {..} = runST $ do+  -- marks both (u, v) and (v, u)+  usedUV <- HM.new @_ @Bit (mCsr + 4)++  -- we'll give unique indices to (u, v) pairs+  let ix u v = min u v .<<. 32 .|. max u v++  -- depth+  dep <- VUM.replicate nCsr (-1 :: Int)++  -- vertex -> edge index+  par <- VUM.replicate nCsr (-1 :: Int)+  parFrom <- VUM.replicate nCsr (-1 :: Int)++  -- Get DFS forest+  let dfs u d = do+        VGM.write dep u d+        VU.forM_ (eAdj gr u) $ \(!iEdge, !v) -> do+          dv <- VGM.read dep v+          when (dv == -1) $ do+            -- we're marking both direction of an undirected edge+            HM.insert usedUV (ix u v) $ Bit True+            VGM.write par v iEdge+            VGM.write parFrom v u+            dfs v (d + 1)++  VGM.iforM_ dep $ \v d -> do+    when (d == -1) $ do+      dfs v 0++  vs <- GV.new @_ @Int 16+  es <- GV.new @_ @Int 16+  dep' <- VU.unsafeFreeze dep++  -- Find edge with minimum depth difference, which makes up a loop (not used in the DFS forets):+  minLen <- VUM.replicate 1 (maxBound `div` 2 :: Int)+  backE <- VUM.replicate 1 (-1 :: Int, -1 :: Int)+  for_ [0 .. nCsr - 1] $ \vA -> do+    let !dA = dep' VG.! vA+    VU.forM_ (eAdj gr vA) $ \(!iEdge, !vB) -> do+      b <- maybe False unBit <$> HM.lookup usedUV (ix vA vB)+      unless b $ do+        let !dB = dep' VG.! vB+        let !d = abs $ dA - dB+        minLen' <- VGM.read minLen 0+        when (d < minLen') $ do+          VGM.write minLen 0 d+          VGM.write backE 0 (iEdge, vA)++  (!backE', !backFrom) <- VGM.read backE 0+  when (backE' /= -1) $ do+    let try a b = do+          if dep' VG.! a > dep' VG.! b+            then try b a+            else do+              -- v_1 -> v_N -> v_{N - 1} -> .. -> v_2 -> v_1+              GV.pushBack es backE'+              GV.pushBack vs a+              let backtrack v = do+                    unless (v == a) $ do+                      parE <- VGM.read par v+                      v' <- VGM.read parFrom v+                      GV.pushBack vs v+                      GV.pushBack es parE+                      backtrack v'+              backtrack b+    try backFrom (adjCsr VG.! backE')++  vs' <- GV.unsafeFreeze vs+  es' <- GV.unsafeFreeze es+  if VU.null es'+    then pure Nothing+    else pure $ Just (vs', es')  -- ------------------------------------------------------------------------------------------------- -- Generic graph search functions
src/AtCoder/Extra/Math/Montgomery64.hs view
@@ -51,15 +51,15 @@     n2M64 :: {-# UNPACK #-} !Word64   }   deriving-    ( -- | @since 1.2.6.0+    ( -- | NOTE: The represented value can be equal even if the internal values are different, so+      -- use `eq` for comparison.+      --+      -- @since 1.2.6.0       Eq,       -- | @since 1.2.6.0       Show     ) --- TODO: add unasfePerformIO?--- TODO: remove NOINLINE?- -- | \(O(1)\) Static, shared storage of `Montgomery64`. -- -- ==== Constraints@@ -67,9 +67,8 @@ -- - \(m\) is odd -- -- @since 1.2.6.0-{-# NOINLINE new #-}+{-# INLINE new #-} new :: forall a. (HasCallStack, KnownNat a) => Proxy# a -> Montgomery64--- FIXME: test allocated once new p = fromVal . fromIntegral $! natVal' p  -- | \(O(1)\) Creates a `Montgomery64` for a modulus value \(m\) of type `Word64` value.
src/AtCoder/Extra/Pdsu.hs view
@@ -34,6 +34,7 @@     -- * Merging     merge,     merge_,+    mergeMaybe,      -- * Group information     size,@@ -69,8 +70,12 @@ -- 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+-- 0 --+-- >>> Pdsu.mergeMaybe dsu 1 0 (Sum 1)+-- Nothing+--+-- -- >>> Pdsu.merge_ dsu 2 0 (Sum 2) -- p(2) - p(0) := Sum 2 -- >>> Pdsu.leader dsu 0 -- 0@@ -172,16 +177,29 @@ -- TODO: use merge and mergeMaybe  -- | \(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.+-- \(\mathrm{dp}\): \(p(v1) := \mathrm{dp} \cdot p(v2)\). Returns the representative of the new+-- connected component. ----- @since 1.1.0.0+-- @since 1.4.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 :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Int merge dsu v10 v20 !dp0 = stToPrim $ mergeST dsu v10 v20 dp0   where     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 $ nPdsu dsu     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 $ nPdsu dsu +-- | \(O(\alpha(n))\) Merges \(v_1\) to \(v_2\) with differential (relative) potential+-- \(\mathrm{dp}\): \(p(v1) := \mathrm{dp} \cdot p(v2)\). Returns the representative of the new+-- connected component if it's newly merged, otherwise it returns `Nothing`.+--+-- @since 1.4.0.0+{-# INLINE mergeMaybe #-}+mergeMaybe :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m (Maybe Int)+mergeMaybe dsu v10 v20 !dp0 = stToPrim $ mergeMaybeST dsu v10 v20 dp0+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.mergeMaybe" v10 $ nPdsu dsu+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.mergeMaybe" v20 $ nPdsu dsu+ -- | \(O(\alpha(n))\) `merge` with the return value discarded. -- -- @since 1.1.0.0@@ -263,14 +281,14 @@     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.potST" v1 nPdsu  {-# INLINEABLE mergeST #-}-mergeST :: (HasCallStack, Monoid a, Ord a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Bool+mergeST :: (HasCallStack, Monoid a, Ord a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Int mergeST dsu@Pdsu {..} v10 v20 !dp0 = inner v10 v20 dp0   where     inner v1 v2 !dp = do       !r1 <- leaderST dsu v1       !r2 <- leaderST dsu v2       if r1 == r2-        then pure False+        then pure r1         else do           -- NOTE(perf): Union by size (choose smaller one for root).           -- Another, more proper optimization would be union by rank (depth).@@ -298,9 +316,18 @@               VGM.write parentOrSizePdsu r1 {- record new root -} r2               VGM.write potentialPdsu r1 pr1' -              pure True+              pure r2             else do               inner v2 v1 $ invertPdsu dp++{-# INLINE mergeMaybeST #-}+mergeMaybeST :: (HasCallStack, Monoid a, Ord a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s (Maybe Int)+mergeMaybeST dsu v1 v2 !dp0 = do+  r1 <- leaderST dsu v1+  r2 <- leaderST dsu v2+  if r1 == r2+    then pure Nothing+    else Just <$> mergeST dsu v1 v2 dp0  {-# INLINEABLE canMergeST #-} canMergeST :: (HasCallStack, Semigroup a, Eq a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Bool
src/AtCoder/Extra/SegTree2d/Dense.hs view
@@ -74,16 +74,12 @@ where  import AtCoder.Internal.Assert qualified as ACIA-import AtCoder.Internal.Bit qualified as ACIB-import Control.Monad (when) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) import Control.Monad.ST (ST) import Data.Bits import Data.Foldable (for_)-import Data.Maybe (fromJust, fromMaybe)+import Data.Maybe (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.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM@@ -164,7 +160,7 @@ -- @since 1.2.3.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m a-read seg@DenseSegTree2d {..} x y = do+read DenseSegTree2d {..} x y = do   let !_ = ACIA.checkPoint2d "AtCoder.Extra.SegTree2d.Dense.read" x y wDst hDst   VGM.read dataDst $ idx wDst (y + hDst) (x + wDst) @@ -174,7 +170,7 @@ -- @since 1.2.3.0 {-# INLINE readMaybe #-} readMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m (Maybe a)-readMaybe seg@DenseSegTree2d {..} x y+readMaybe DenseSegTree2d {..} x y   | ACIA.testPoint2d x y wDst hDst = do       Just <$> VGM.read dataDst (idx wDst (y + hDst) (x + wDst))   | otherwise = pure Nothing@@ -245,7 +241,6 @@ {-# INLINE allProd #-} allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> m a allProd DenseSegTree2d {..} = stToPrim $ do-  -- FIXME: correct?   fromMaybe mempty <$> VGM.readMaybe dataDst (idx wDst 1 1)  -- -------------------------------------------------------------------------------------------------
src/AtCoder/Extra/Seq/Map.hs view
@@ -59,7 +59,7 @@     -- sliceST,     prod,     prodMaybe,-    allProd, -- FIXME: rename to `prodAll`+    allProd,      -- ** Applications     applyIn,
src/AtCoder/Extra/Tree/Hld.hs view
@@ -376,12 +376,13 @@         ihv = indexHld VG.! hv  -- | \(O(\log n)\) Returns the \(k\)-th vertex of the path between \(u\) and \(v\) from \(u\).--- Returns `Nothing` if \(k\) is bigger than the path length.+-- Returns `Nothing` if \(k\) is bigger than the path length or \(k < 0\). ----- @since 1.1.0.0+-- @since 1.4.0.0 {-# INLINEABLE jump #-} jump :: (HasCallStack) => Hld -> Vertex -> Vertex -> Int -> Maybe Vertex jump hld@Hld {..} u v k+  | k < 0 = Nothing   | k > lenU + lenV = Nothing   | k <= lenU = Just $ ancestor hld u k   | otherwise = Just $ ancestor hld v (lenU + lenV - k)
src/AtCoder/Extra/WaveletMatrix2d.hs view
@@ -2,7 +2,7 @@  -- | 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.+-- modified later. Duplicate monoids at the same coordinate will be combined into one. -- -- ==== `SegTree2d` vs `WaveletMatrix2d` -- They basically have the same functionalities and performance, however, `SegTree2d` performs better in@@ -55,7 +55,7 @@     modify,     prod,     prodMaybe,-    allProd, -- FIXME: rename to prodAll+    allProd,     -- wavelet matrix methods could be implemented, too   ) where@@ -133,8 +133,8 @@   segTreesWm2d <- V.replicateM (Rwm.heightRwm rawWm2d) (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 `(<>)`.+-- | \(O(n \log n)\) Creates a `WaveletMatrix2d` with wavelet matrix with segment tree with initial+-- monoid values. Duplicate monoids at the same coordinate will be combined with `(<>)`. -- -- @since 1.1.0.0 {-# INLINEABLE build #-}
src/AtCoder/Internal/Csr.hs view
@@ -44,6 +44,7 @@     adj,     adjW,     eAdj,+    eAdjW,   ) where @@ -67,7 +68,7 @@     --     -- @since 1.1.0.0     mCsr :: {-# UNPACK #-} !Int,-    -- | Starting indices.+    -- | Maps vertices into the starting indices for `adjCsr` and `wCsr`.     --     -- @since 1.1.0.0     startCsr :: !(VU.Vector Int),@@ -157,7 +158,7 @@       ir = startCsr VG.! (i + 1)    in VU.zip (VU.slice il (ir - il) adjCsr) (VU.slice il (ir - il) wCsr) --- | \(O(n)\) Returns a vector of @(edgeId, adjacentVertex)@.+-- | \(O(n)\) Returns a vector of @(csrEdgeIndex, adjacentVertex)@. -- -- @since 1.0.0.0 {-# INLINE eAdj #-}@@ -166,3 +167,13 @@   let il = startCsr VG.! i       ir = startCsr VG.! (i + 1)    in VU.imap ((,) . (+ il)) $ VU.slice il (ir - il) adjCsr++-- | \(O(n)\) Returns a vector of @(csrEdgeIndex, adjacentVertex, edgeWeight)@.+--+-- @since 1.4.0.0+{-# INLINE eAdjW #-}+eAdjW :: (HasCallStack, VU.Unbox w) => Csr w -> Int -> VU.Vector (Int, Int, w)+eAdjW Csr {..} i =+  let il = startCsr VG.! i+      ir = startCsr VG.! (i + 1)+   in VU.zip3 (VU.enumFromN il (ir - il)) (VU.slice il (ir - il) adjCsr) (VU.slice il (ir - il) wCsr)
src/AtCoder/Internal/GrowVec.hs view
@@ -64,6 +64,8 @@     -- * Readings     read,     readMaybe,+    readBack,+    readBackMaybe,      -- * Modifications @@ -75,12 +77,18 @@     popBack,     popBack_, +    -- ** Misc+    clear,+    reverse,+     -- * Conversion     freeze,     unsafeFreeze,   ) where +-- NOTE (perf): we have to inine all the functions for reasonable performance+ import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (when) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)@@ -90,13 +98,15 @@ import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM import GHC.Stack (HasCallStack)-import Prelude hiding (length, null, read)+import Prelude hiding (length, null, read, reverse)  -- | Growable vector with some runtime overhead. -- -- @since 1.0.0.0 data GrowVec s a = GrowVec   { -- | Stores [l, r) range in the `vecGV`.+    --+    -- @since 1.0.0.0     posGV :: !(VUM.MVector s Int),     -- | @since 1.0.0.0     vecGV :: !(MutVar s (VUM.MVector s a))@@ -160,7 +170,7 @@ -- | \(O(1)\) Yields the element at the given position. Will throw an exception if the index is out -- of range. ----- @since 1.0.0.0+-- @since 1.4.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m a read gv i = stToPrim $ readST gv i@@ -172,10 +182,25 @@ readMaybe :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m (Maybe a) readMaybe gv i = stToPrim $ readMaybeST gv i +-- | \(O(1)\) Yields the element at the given position. Will throw an exception if the index is out+-- of range.+--+-- @since 1.4.0.0+{-# INLINE readBack #-}+readBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m a+readBack gv i = stToPrim $ readBackST gv i++-- | \(O(1)\) Yields the element at the given position, or `Nothing` if the index is out of range.+--+-- @since 1.4.0.0+{-# INLINE readBackMaybe #-}+readBackMaybe :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m (Maybe a)+readBackMaybe gv i = stToPrim $ readBackMaybeST gv i+ -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of range. ----- @since 1.0.0.0+-- @since 1.4.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> a -> m () write gv i x = stToPrim $ writeST gv i x@@ -201,6 +226,21 @@ popBack_ :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m () popBack_ = stToPrim . popBackST_ +-- | \(O(1)\) Sets the length to zero.+--+-- @since 1.4.0.0+{-# INLINE clear #-}+clear :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m ()+clear GrowVec {..} = stToPrim $ do+  VGM.write posGV 0 0++-- | \(O(n)\) Reverses all the elements.+--+-- @since 1.4.0.0+{-# INLINE reverse #-}+reverse :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => GrowVec (PrimState m) a -> m ()+reverse = stToPrim . reverseST+ -- | \(O(n)\) Yields an immutable copy of the mutable vector. -- -- @since 1.0.0.0@@ -220,15 +260,15 @@ -- Internal -- ------------------------------------------------------------------------------------------------- -{-# INLINEABLE readST #-}+{-# INLINE readST #-} readST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s a readST GrowVec {..} i = do   vec <- readMutVar vecGV-  let len = VUM.length vec-  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.read" i len+  len <- VGM.unsafeRead posGV 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.readST" i len   VGM.read vec i -{-# INLINEABLE readMaybeST #-}+{-# INLINE readMaybeST #-} readMaybeST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s (Maybe a) readMaybeST GrowVec {..} i = do   vec <- readMutVar vecGV@@ -237,15 +277,32 @@     then Just <$> VGM.unsafeRead vec i     else pure Nothing -{-# INLINEABLE writeST #-}+{-# INLINE readBackST #-}+readBackST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s a+readBackST GrowVec {..} i = do+  vec <- readMutVar vecGV+  len <- VGM.unsafeRead posGV 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.readBackST" i len+  VGM.read vec (len - 1 - i)++{-# INLINE readBackMaybeST #-}+readBackMaybeST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s (Maybe a)+readBackMaybeST GrowVec {..} i = do+  vec <- readMutVar vecGV+  len <- VGM.unsafeRead posGV 0+  if ACIA.testIndex i len+    then Just <$> VGM.unsafeRead vec (len - 1 - i)+    else pure Nothing++{-# INLINE writeST #-} writeST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> a -> ST s () writeST GrowVec {..} i x = do   vec <- readMutVar vecGV-  let len = VUM.length vec-  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.write" i len+  len <- VGM.unsafeRead posGV 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.writeST" i len   VGM.write vec i x -{-# INLINEABLE pushBackST #-}+{-# INLINE pushBackST #-} pushBackST :: (VU.Unbox a) => GrowVec s a -> a -> ST s () pushBackST GrowVec {..} e = do   len <- VGM.unsafeRead posGV 0@@ -267,7 +324,7 @@     )     0 -{-# INLINEABLE popBackST #-}+{-# INLINE popBackST #-} popBackST :: (VU.Unbox a) => GrowVec s a -> ST s (Maybe a) popBackST GrowVec {..} = do   pos <- VGM.unsafeRead posGV 0@@ -278,20 +335,28 @@       vec <- readMutVar vecGV       Just <$> VGM.read vec (pos - 1) -{-# INLINEABLE popBackST_ #-}+{-# INLINE popBackST_ #-} popBackST_ :: (VU.Unbox a) => GrowVec s a -> ST s () popBackST_ GrowVec {..} = do   pos <- VGM.unsafeRead posGV 0   VGM.unsafeWrite posGV 0 $ max 0 $ pos - 1 -{-# INLINEABLE freezeST #-}+{-# INLINE reverseST #-}+reverseST :: (HasCallStack, Ord a, VU.Unbox a) => GrowVec s a -> ST s ()+reverseST GrowVec {..} = do+  len <- VGM.unsafeRead posGV 0+  vec <- readMutVar vecGV+  let slice = VUM.take len vec+  VGM.reverse slice++{-# INLINE freezeST #-} freezeST :: (VU.Unbox a) => GrowVec s a -> ST s (VU.Vector a) freezeST GrowVec {..} = do   len <- VGM.unsafeRead posGV 0   vec <- readMutVar vecGV   VU.freeze $ VUM.take len vec -{-# INLINEABLE unsafeFreezeST #-}+{-# INLINE unsafeFreezeST #-} unsafeFreezeST :: (VU.Unbox a) => GrowVec s a -> ST s (VU.Vector a) unsafeFreezeST GrowVec {..} = do   len <- VGM.unsafeRead posGV 0
src/AtCoder/LazySegTree.hs view
@@ -97,7 +97,7 @@     -- * Products     prod,     prodMaybe,-    allProd, -- FIXME: rename to prodAll+    allProd,      -- * Applications     applyAt,
test/Main.hs view
@@ -24,6 +24,7 @@ import Tests.Extra.ModInt64 qualified import Tests.Extra.Monoid qualified import Tests.Extra.MultiSet qualified+import Tests.Extra.Pdsu qualified import Tests.Extra.SegTree2d qualified import Tests.Extra.SegTree2d.Dense qualified import Tests.Extra.Semigroup.Matrix qualified@@ -81,6 +82,7 @@             testGroup "ModInt64" Tests.Extra.ModInt64.tests,             testGroup "Monoid" Tests.Extra.Monoid.tests,             testGroup "MultiSet" Tests.Extra.MultiSet.tests,+            testGroup "Pdsu" Tests.Extra.Pdsu.tests,             testGroup "SegTree2d" Tests.Extra.SegTree2d.tests,             testGroup "SegTree2d.Dense" Tests.Extra.SegTree2d.Dense.tests,             testGroup "Semigroup.Matrix" Tests.Extra.Semigroup.Matrix.tests,
+ test/Tests/Extra/Pdsu.hs view
@@ -0,0 +1,51 @@+{-# LANGUAGE DerivingStrategies #-}++module Tests.Extra.Pdsu (tests) where++import AtCoder.Extra.Pdsu qualified as Pdsu+import Control.Monad (forM)+import Control.Monad.ST (runST)+import Data.Maybe (fromJust)+import Test.Tasty+import Test.Tasty.QuickCheck as QC++prop_merge :: QC.Positive Int -> QC.Gen QC.Property+prop_merge (QC.Positive n) = do+  m <- QC.chooseInt (1, 2 * n)+  es <- QC.vectorOf m $ do+    u <- QC.chooseInt (0, n - 1)+    v <- QC.chooseInt (0, n - 1)+    pure (u, v)+  pure . QC.conjoin $ runST $ do+    dsu <- Pdsu.new @_ @() n id+    forM es $ \(!u, !v) -> do+      r' <- Pdsu.merge dsu u v ()+      r1 <- Pdsu.leader dsu u+      r2 <- Pdsu.leader dsu v+      pure (r1 == r' && r2 == r')++prop_mergeMaybe :: QC.Positive Int -> QC.Gen QC.Property+prop_mergeMaybe (QC.Positive n) = do+  m <- QC.chooseInt (1, 2 * n)+  es <- QC.vectorOf m $ do+    u <- QC.chooseInt (0, n - 1)+    v <- QC.chooseInt (0, n - 1)+    pure (u, v)+  pure . QC.conjoin $ runST $ do+    dsu <- Pdsu.new @_ @() n id+    forM es $ \(!u, !v) -> do+      r1 <- Pdsu.leader dsu u+      r2 <- Pdsu.leader dsu v+      r' <- Pdsu.mergeMaybe dsu u v ()+      if r1 == r2+        then pure $ r' == Nothing+        else do+          r1' <- Pdsu.leader dsu u+          r2' <- Pdsu.leader dsu v+          pure (r1' == fromJust r' && r2' == fromJust r')++tests :: [TestTree]+tests =+  [ QC.testProperty "merge" prop_merge,+    QC.testProperty "mergeMaybe" prop_mergeMaybe+  ]