ac-library-hs-1.2.3.0: src/AtCoder/Extra/LazyKdTree.hs
{-# LANGUAGE RecordWildCards #-}
-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
-- commutative monoids.
--
-- - Point coordinates are fixed on `build`.
-- - Multiple points can exist at the same coordinate.
--
-- ==== __Examples__
-- >>> import AtCoder.Extra.LazyKdTree qualified as LKT
-- >>> import AtCoder.Extra.Monoid.Affine1 (Affine1)
-- >>> import AtCoder.Extra.Monoid.Affine1 qualified as Affine1
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let xyws = VU.fromList [(0, 0, Sum 1), (1, 1, Sum 2), (4, 2, Sum 3)]
-- >>> lkt <- LKT.build3 @_ @(Affine1 Int) @(Sum Int) xyws
--
-- >>> -- Get monoid product in [0, 2) x [0, 2)
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 3}
--
-- >>> LKT.applyIn lkt 0 2 0 2 $ Affine1.new 2 1
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 8}
--
-- >>> LKT.write lkt 0 $ Sum 10
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 15}
--
-- @since 1.2.2.0
module AtCoder.Extra.LazyKdTree
( -- * K-dimensional tree
LazyKdTree (..),
-- * Re-exports
SegAct (..),
-- * Constructors
new,
build,
build2,
build3,
-- * Write
write,
modify,
modifyM,
-- * Monoid products
prod,
allProd,
-- * Apply
applyIn,
)
where
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Bit qualified as ACIB
import AtCoder.LazySegTree (SegAct (..))
import Control.Monad (unless, when)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bits
import Data.Foldable (for_)
import Data.Maybe (fromMaybe)
import Data.Ord (comparing)
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
import GHC.Stack (HasCallStack)
-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
-- commutative monoids.
--
-- @since 1.2.2.0
data LazyKdTree s f a = LazyKdTree
{ -- | The number of points in the \(k\)-d tree.
--
-- @since 1.2.2.0
nLkt :: {-# UNPACK #-} !Int,
-- | \(\lceil \log_2 (n + 1) \rceil\)
--
-- @since 1.2.2.0
logLkt :: {-# UNPACK #-} !Int,
-- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\).
--
-- @since 1.2.2.0
incRectsLkt :: !(VU.Vector (Int, Int, Int, Int)),
-- | Rectangle information: monoid values.
--
-- @since 1.2.2.0
dataLkt :: !(VUM.MVector s a),
-- | Rectangle information: lazily propagated monoid actions for children.
--
-- @since 1.2.2.0
lazyLkt :: !(VUM.MVector s f),
-- | Rectangle information: the number of vertices in the rectangle.
--
-- @since 1.2.2.0
sizeLkt :: !(VU.Vector Int),
-- | Maps original vertices into the belonging rectangle index.
--
-- @since 1.2.2.0
posLkt :: !(VU.Vector Int)
}
-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@ and @ys@.
--
-- ==== Constraints
-- - \(|\mathrm{xs}| = |\mathrm{ys}|
--
-- @since 1.2.3.0
{-# INLINE new #-}
new ::
(HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | \(x\) coordnates
VU.Vector Int ->
-- | \(y\) coordnates
VU.Vector Int ->
-- | `LazyKdTree`
m (LazyKdTree (PrimState m) f a)
new xs ys = stToPrim $ buildST xs ys (VU.replicate (VU.length xs) mempty)
-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@, @ys@ and @ws@ vectors.
--
-- ==== Constraints
-- - \(|\mathrm{xs}| = |\mathrm{ys}| = |\mathrm{vs}|\).
--
-- @since 1.2.2.0
{-# INLINE build #-}
build ::
(HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | \(x\) coordnates
VU.Vector Int ->
-- | \(y\) coordnates
VU.Vector Int ->
-- | monoid \(v\)alues
VU.Vector a ->
-- | `LazyKdTree`
m (LazyKdTree (PrimState m) f a)
build xs ys vs = stToPrim $ buildST xs ys vs
-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xys@ and @ws@ vectors.
--
-- ==== Constraints
-- - \(|\mathrm{xys}| = |\mathrm{vs}|\).
--
-- @since 1.2.2.0
{-# INLINE build2 #-}
build2 ::
(HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | \((x, y)\) coordinates
VU.Vector (Int, Int) ->
-- | Monoid \(v\)alues
VU.Vector a ->
-- | `LazyKdTree`
m (LazyKdTree (PrimState m) f a)
build2 xys ws = stToPrim $ buildST xs ys ws
where
(!xs, !ys) = VU.unzip xys
-- | \(O(n \log n)\) Creates a `LazyKdTree` from a @xyws@ vector.
--
-- @since 1.2.2.0
{-# INLINE build3 #-}
build3 ::
(HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | \((x, y, v)\) tuples
VU.Vector (Int, Int, a) ->
-- | `LazyKdTree`
m (LazyKdTree (PrimState m) f a)
build3 xyws = stToPrim $ buildST xs ys ws
where
(!xs, !ys, !ws) = VU.unzip3 xyws
-- | \(O(\log n)\) Writes to the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINE write #-}
write ::
(HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | Original vertex index.
Int ->
-- | Monoid value
a ->
-- | Monadic tuple
m ()
write kt i x = stToPrim $ modifyM kt (pure . const x) i
-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINE modify #-}
modify ::
(HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | Creates a new monoid value from the old one.
(a -> a) ->
-- | Original vertex index.
Int ->
-- | Monadic tuple
m ()
modify kt f i = stToPrim $ modifyM kt (pure . f) i
-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINEABLE modifyM #-}
modifyM ::
(HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | Creates a new monoid value from the old one.
(a -> m a) ->
-- | Original vertex index.
Int ->
-- | Monadic tuple
m ()
modifyM kt@LazyKdTree {..} f i0 = do
let !_ = ACIA.checkIndex "AtCoder.Extra.LazyKdTree.modifyM" i0 nLkt
let i_ = posLkt VG.! i0
-- propagate lazily propagated monoid actions from the root:
stToPrim $ for_ [logLkt, logLkt - 1 .. 1] $ \k -> do
pushST kt (i_ .>>. k)
VGM.modifyM dataLkt f i_
-- update parents:
let inner i
| i <= 1 = pure ()
| otherwise = do
let i' = i `div` 2
xl <- VGM.read dataLkt (2 * i' + 0)
xr <- VGM.read dataLkt (2 * i' + 1)
VGM.write dataLkt i' $! xl <> xr
inner i'
stToPrim $ inner i_
-- | \(O(\sqrt n)\) Returns monoid product \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\).
--
-- @since 1.2.2.0
{-# INLINE prod #-}
prod ::
(HasCallStack, PrimMonad m, Eq f, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | \(x_1\)
Int ->
-- | \(x_2\)
Int ->
-- | \(y_1\)
Int ->
-- | \(y_2\)
Int ->
-- | \(\Pi_{p \in [x_1, x_2) \times [y_1, y_2)} a_p\)
m a
prod kt x1 x2 y1 y2 = stToPrim $ prodST kt x1 x2 y1 y2
-- | \(O(1)\) Returns monoid product \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\).
--
-- @since 1.2.2.0
{-# INLINE allProd #-}
allProd ::
(PrimMonad m, Monoid a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | \(\Pi_{p \in [-\infty, \infty) \times [-\infty, \infty)} a_p\)
m a
allProd kt = do
-- In case of zero vertices, use `Maybe`:
fromMaybe mempty <$> VGM.readMaybe (dataLkt kt) 1
-- | \(O(\log n)\) Applies a monoid action to points in \([x_1, x_2) \times [y_1, y_2)\).
--
-- @since 1.2.2.0
{-# INLINE applyIn #-}
applyIn ::
(HasCallStack, PrimMonad m, Eq f, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | `LazyKdTree`
LazyKdTree (PrimState m) f a ->
-- | \(x_1\)
Int ->
-- | \(x_2\)
Int ->
-- | \(y_1\)
Int ->
-- | \(y_2\)
Int ->
-- | \(f\)
f ->
-- | Monadic tuple
m ()
applyIn kt x1 x2 y1 y2 f = stToPrim $ applyInST kt 1 x1 x2 y1 y2 f
-- -------------------------------------------------------------------------------------------------
-- Private
-- -------------------------------------------------------------------------------------------------
{-# INLINEABLE buildST #-}
buildST :: forall s f a. (HasCallStack, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (LazyKdTree s f a)
buildST xs0 ys0 vs0 = do
let nLkt = VU.length xs0
let !_ = ACIA.runtimeAssert (nLkt == VU.length ys0 && nLkt == VU.length vs0) "AtCoder.Extra.LazyKdTree.buildST: the length of `xs`, `ys` and `vs` must be equal"
if nLkt == 0
then do
let logLkt = 0
dataLkt <- VUM.new 0
lazyLkt <- VUM.new 0
let incRectsLkt = VU.empty
let sizeLkt = VU.empty
let posLkt = VU.empty
pure LazyKdTree {..}
else do
let logLkt = countTrailingZeros $ ACIB.bitCeil (nLkt + 1)
dataLkt <- VUM.unsafeNew (bit (logLkt + 1))
lazyLkt <- VUM.replicate (bit logLkt) mempty
incRectsVec <- VUM.replicate (bit (logLkt + 1)) (maxBound, minBound, maxBound, minBound)
size <- VUM.unsafeNew (bit (logLkt + 1))
pos <- VUM.unsafeNew nLkt
let VUM.MV_4 _ xMins xMaxes yMins yMaxes = incRectsVec
-- - idx: rectangle index (one-based)
-- - xs, ys, vs: point information (x, y and monoid value)
-- - ids: maps sorted vertices to the original vertex indices
-- - divX: represents hyperplane direction for point partition
let buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector a -> VU.Vector Int -> Bool -> ST s ()
buildSubtree idx xs ys vs ids divX = do
let n = VU.length xs
VGM.write size idx n
-- retrieve the bounds:
let (!xMin, !xMax, !yMin, !yMax) =
VU.foldl'
(\(!a, !b, !c, !d) (!x, !y) -> (min a x, max b x, min c y, max d y))
(maxBound, minBound, maxBound, minBound)
$ VU.zip xs ys
VGM.modify xMins (min xMin) idx
VGM.modify xMaxes (max xMax) idx
VGM.modify yMins (min yMin) idx
VGM.modify yMaxes (max yMax) idx
if n == 1
then do
-- it's a terminal. note that it's not always a leaf; the case is handled carefully in
-- other methods
VGM.write dataLkt idx $ vs VG.! 0
-- record original vertex index -> rectangle index
VGM.write pos (ids VG.! 0) idx
else do
-- partition the vertices into two:
let m = n `div` 2
let is = VU.create $ do
vec <- VUM.generate n id
if divX
then VAI.selectBy (comparing (xs VG.!)) vec m
else VAI.selectBy (comparing (ys VG.!)) vec m
pure vec
-- TODO: permute in-place?
let (!xsL, !xsR) = VG.splitAt m $ VG.backpermute xs is
let (!ysL, !ysR) = VG.splitAt m $ VG.backpermute ys is
let (!vsL, !vsR) = VG.splitAt m $ VG.backpermute vs is
let (!idsL, !idsR) = VG.splitAt m $ VG.backpermute ids is
-- build the subtree:
buildSubtree (2 * idx + 0) xsL ysL vsL idsL (not divX)
buildSubtree (2 * idx + 1) xsR ysR vsR idsR (not divX)
xl <- VGM.read dataLkt (2 * idx + 0)
xr <- VGM.read dataLkt (2 * idx + 1)
VGM.write dataLkt idx $! xl <> xr
buildSubtree 1 xs0 ys0 vs0 (VU.generate nLkt id) True
sizeLkt <- VU.unsafeFreeze size
posLkt <- VU.unsafeFreeze pos
incRectsLkt <- VU.unsafeFreeze incRectsVec
pure LazyKdTree {..}
{-# INLINE applyAtST #-}
applyAtST :: (SegAct f a, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST LazyKdTree {..} i f = do
-- NOTE: Here we're asssuming each monoid value has length one. If you need a monoid of length
-- zero, e.g., if you're just reserving new point insertion, you must not rely on
-- `segActWithLength`. You might want to use `V2` instead of `Sum`.
let len = sizeLkt VG.! i
VGM.modify dataLkt (segActWithLength len f) i
when (i < bit logLkt) $ do
VGM.modify lazyLkt (f <>) i
-- TODO: consider `INLINE`?
{-# INLINE pushST #-}
pushST :: (SegAct f a, Eq f, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> ST s ()
pushST kt@LazyKdTree {..} i = do
lazy <- VGM.read lazyLkt i
unless (lazy == mempty) $ do
applyAtST kt (2 * i + 0) lazy
applyAtST kt (2 * i + 1) lazy
VGM.write lazyLkt i mempty
{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a
prodST kt@LazyKdTree {..} x1 x2 y1 y2
| x1 >= x2 || y1 >= y2 = pure mempty
| otherwise = inner 1
where
inner i = case incRectsLkt VG.!? i of
Nothing -> pure mempty
Just (!xl, !xr, !yl, !yr)
-- TODO: what is this?
| xl > xr -> pure mempty
-- not intersecting
| x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure mempty
-- the rectangle is fully contained by the query:
| x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do
VGM.read dataLkt i
| otherwise -> do
pushST kt i
l <- inner (2 * i + 0)
r <- inner (2 * i + 1)
pure $! l <> r
{-# INLINEABLE applyInST #-}
applyInST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s ()
applyInST kt@LazyKdTree {..} i0 x1 x2 y1 y2 f
| x1 >= x2 || y1 >= y2 = pure ()
| otherwise = inner i0
where
inner i = case incRectsLkt VG.!? i of
Nothing -> pure mempty
Just (!xl, !xr, !yl, !yr)
| xl > xr -> pure ()
-- not intersecting
| x2 <= xl || x1 > xr || y2 <= yl || y1 > yr -> pure ()
-- the rectangle is fully contained by the query:
| x1 <= xl && xr < x2 && y1 <= yl && yr < y2 -> do
applyAtST kt i f
| otherwise -> do
pushST kt i
inner (2 * i + 0)
inner (2 * i + 1)
l <- VGM.read dataLkt (2 * i + 0)
r <- VGM.read dataLkt (2 * i + 1)
VGM.write dataLkt i $! l <> r