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 +11/−0
- ac-library-hs.cabal +5/−1
- benchmarks/Bench/Montgomery64.hs +41/−0
- benchmarks/BenchLib/Montgomery64/Inline.hs +40/−0
- benchmarks/BenchLib/Montgomery64/Noinline.hs +40/−0
- benchmarks/Main.hs +2/−3
- src/AtCoder/Dsu.hs +3/−0
- src/AtCoder/Extra/Bisect.hs +6/−6
- src/AtCoder/Extra/DynLazySegTree.hs +1/−1
- src/AtCoder/Extra/DynLazySegTree/Persistent.hs +1/−1
- src/AtCoder/Extra/DynSegTree.hs +1/−1
- src/AtCoder/Extra/DynSegTree/Persistent.hs +1/−1
- src/AtCoder/Extra/DynSparseSegTree.hs +1/−1
- src/AtCoder/Extra/DynSparseSegTree/Persistent.hs +1/−1
- src/AtCoder/Extra/Graph.hs +286/−3
- src/AtCoder/Extra/Math/Montgomery64.hs +5/−6
- src/AtCoder/Extra/Pdsu.hs +34/−7
- src/AtCoder/Extra/SegTree2d/Dense.hs +3/−8
- src/AtCoder/Extra/Seq/Map.hs +1/−1
- src/AtCoder/Extra/Tree/Hld.hs +3/−2
- src/AtCoder/Extra/WaveletMatrix2d.hs +4/−4
- src/AtCoder/Internal/Csr.hs +13/−2
- src/AtCoder/Internal/GrowVec.hs +80/−15
- src/AtCoder/LazySegTree.hs +1/−1
- test/Main.hs +2/−0
- test/Tests/Extra/Pdsu.hs +51/−0
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+ ]