{-# LANGUAGE MultiParamTypeClasses #-}
module SegTreeLazy
( LazySegTree
, emptyLST
, fromListLST
, adjustLST
, updateRangeLST
, foldRangeLST
, foldrLST
) where
import Control.DeepSeq
import Control.Monad.State
import Data.Bits
import Misc ( Action(..), bitLength )
data LazySegTree u a = LazySegTree !(Int, Int, Int) !(LSegNode u a) deriving Int -> LazySegTree u a -> ShowS
[LazySegTree u a] -> ShowS
LazySegTree u a -> String
(Int -> LazySegTree u a -> ShowS)
-> (LazySegTree u a -> String)
-> ([LazySegTree u a] -> ShowS)
-> Show (LazySegTree u a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall u a. (Show a, Show u) => Int -> LazySegTree u a -> ShowS
forall u a. (Show a, Show u) => [LazySegTree u a] -> ShowS
forall u a. (Show a, Show u) => LazySegTree u a -> String
showList :: [LazySegTree u a] -> ShowS
$cshowList :: forall u a. (Show a, Show u) => [LazySegTree u a] -> ShowS
show :: LazySegTree u a -> String
$cshow :: forall u a. (Show a, Show u) => LazySegTree u a -> String
showsPrec :: Int -> LazySegTree u a -> ShowS
$cshowsPrec :: forall u a. (Show a, Show u) => Int -> LazySegTree u a -> ShowS
Show
data LSegNode u a = LSLeaf !a | LSBin !a !u !(LSegNode u a) !(LSegNode u a) deriving Int -> LSegNode u a -> ShowS
[LSegNode u a] -> ShowS
LSegNode u a -> String
(Int -> LSegNode u a -> ShowS)
-> (LSegNode u a -> String)
-> ([LSegNode u a] -> ShowS)
-> Show (LSegNode u a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall u a. (Show a, Show u) => Int -> LSegNode u a -> ShowS
forall u a. (Show a, Show u) => [LSegNode u a] -> ShowS
forall u a. (Show a, Show u) => LSegNode u a -> String
showList :: [LSegNode u a] -> ShowS
$cshowList :: forall u a. (Show a, Show u) => [LSegNode u a] -> ShowS
show :: LSegNode u a -> String
$cshow :: forall u a. (Show a, Show u) => LSegNode u a -> String
showsPrec :: Int -> LSegNode u a -> ShowS
$cshowsPrec :: forall u a. (Show a, Show u) => Int -> LSegNode u a -> ShowS
Show
buildLST :: Action u a => (Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
buildLST :: (Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
buildLST (Int
l, Int
r) Int -> LSegNode u a
f
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< -Int
1 = String -> LazySegTree u a
forall a. HasCallStack => String -> a
error String
"invalid range"
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1 = (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
forall u a. (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
LazySegTree (Int
l, Int
r, Int
0) (a -> LSegNode u a
forall u a. a -> LSegNode u a
LSLeaf a
forall a. Monoid a => a
mempty)
| Bool
otherwise = (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
forall u a. (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
LazySegTree (Int
l, Int
r, Int -> Int
forall a. Bits a => Int -> a
bit Int
ht) (Int -> LSegNode u a
f Int
ht)
where
n :: Int
n = Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
ht :: Int
ht = Int -> Int
forall b. FiniteBits b => b -> Int
bitLength Int
n
emptyLST :: Action u a => (Int, Int) -> LazySegTree u a
emptyLST :: (Int, Int) -> LazySegTree u a
emptyLST (Int, Int)
bnds = (Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
forall u a.
Action u a =>
(Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
buildLST (Int, Int)
bnds Int -> LSegNode u a
forall a a u.
(Eq a, Monoid a, Monoid u, Num a) =>
a -> LSegNode u a
go where
go :: a -> LSegNode u a
go a
0 = a -> LSegNode u a
forall u a. a -> LSegNode u a
LSLeaf a
forall a. Monoid a => a
mempty
go a
j = a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
forall u a. a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
LSBin a
forall a. Monoid a => a
mempty u
forall a. Monoid a => a
mempty LSegNode u a
lr LSegNode u a
lr where lr :: LSegNode u a
lr = a -> LSegNode u a
go (a
j a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
makeLSN :: Action u a => LSegNode u a -> LSegNode u a -> LSegNode u a
makeLSN :: LSegNode u a -> LSegNode u a -> LSegNode u a
makeLSN LSegNode u a
lt LSegNode u a
rt = a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
forall u a. a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
LSBin (LSegNode u a -> a
forall u p. LSegNode u p -> p
getx LSegNode u a
lt a -> a -> a
forall a. Semigroup a => a -> a -> a
<> LSegNode u a -> a
forall u p. LSegNode u p -> p
getx LSegNode u a
rt) u
forall a. Monoid a => a
mempty LSegNode u a
lt LSegNode u a
rt where
getx :: LSegNode u p -> p
getx (LSLeaf p
x) = p
x
getx (LSBin p
x u
_ LSegNode u p
_ LSegNode u p
_) = p
x
fromListLST :: Action u a => (Int, Int) -> [a] -> LazySegTree u a
fromListLST :: (Int, Int) -> [a] -> LazySegTree u a
fromListLST (Int, Int)
bnds [a]
xs = (Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
forall u a.
Action u a =>
(Int, Int) -> (Int -> LSegNode u a) -> LazySegTree u a
buildLST (Int, Int)
bnds ((State [a] (LSegNode u a) -> [a] -> LSegNode u a)
-> [a] -> State [a] (LSegNode u a) -> LSegNode u a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State [a] (LSegNode u a) -> [a] -> LSegNode u a
forall s a. State s a -> s -> a
evalState [a]
xs (State [a] (LSegNode u a) -> LSegNode u a)
-> (Int -> State [a] (LSegNode u a)) -> Int -> LSegNode u a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> State [a] (LSegNode u a)
forall a u.
(Eq a, Num a, Action u a) =>
a -> StateT [a] Identity (LSegNode u a)
go) where
pop :: StateT [a] Identity a
pop = ([a] -> (a, [a])) -> StateT [a] Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state [a] -> (a, [a])
forall a. Monoid a => [a] -> (a, [a])
go where
go :: [a] -> (a, [a])
go [] = (a
forall a. Monoid a => a
mempty, [])
go (a
x:[a]
xs) = (a
x, [a]
xs)
go :: a -> StateT [a] Identity (LSegNode u a)
go a
0 = a -> LSegNode u a
forall u a. a -> LSegNode u a
LSLeaf (a -> LSegNode u a)
-> StateT [a] Identity a -> StateT [a] Identity (LSegNode u a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT [a] Identity a
pop
go a
j = LSegNode u a -> LSegNode u a -> LSegNode u a
forall u a.
Action u a =>
LSegNode u a -> LSegNode u a -> LSegNode u a
makeLSN (LSegNode u a -> LSegNode u a -> LSegNode u a)
-> StateT [a] Identity (LSegNode u a)
-> StateT [a] Identity (LSegNode u a -> LSegNode u a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> StateT [a] Identity (LSegNode u a)
go (a
j a -> a -> a
forall a. Num a => a -> a -> a
- a
1) StateT [a] Identity (LSegNode u a -> LSegNode u a)
-> StateT [a] Identity (LSegNode u a)
-> StateT [a] Identity (LSegNode u a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> StateT [a] Identity (LSegNode u a)
go (a
j a -> a -> a
forall a. Num a => a -> a -> a
- a
1)
applyLSN :: Action u a => LSegNode u a -> u -> LSegNode u a
applyLSN :: LSegNode u a -> u -> LSegNode u a
applyLSN (LSLeaf a
x) u
u' = a -> LSegNode u a
forall u a. a -> LSegNode u a
LSLeaf (a -> u -> a
forall u a. Action u a => a -> u -> a
act a
x u
u')
applyLSN (LSBin a
x u
u LSegNode u a
lt LSegNode u a
rt) u
u' = a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
forall u a. a -> u -> LSegNode u a -> LSegNode u a -> LSegNode u a
LSBin (a -> u -> a
forall u a. Action u a => a -> u -> a
act a
x u
u') (u
u u -> u -> u
forall a. Semigroup a => a -> a -> a
<> u
u') LSegNode u a
lt LSegNode u a
rt
adjustLST :: Action u a => (a -> a) -> Int -> LazySegTree u a -> LazySegTree u a
adjustLST :: (a -> a) -> Int -> LazySegTree u a -> LazySegTree u a
adjustLST a -> a
f Int
i (LazySegTree lrp :: (Int, Int, Int)
lrp@(Int
l0,Int
r0,Int
p) LSegNode u a
root)
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l0 Bool -> Bool -> Bool
|| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i = String -> LazySegTree u a
forall a. HasCallStack => String -> a
error String
"adjustLST: outside range"
| Bool
otherwise = (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
forall u a. (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
LazySegTree (Int, Int, Int)
lrp (LSegNode u a -> Int -> Int -> u -> LSegNode u a
forall t.
Action t a =>
LSegNode t a -> Int -> Int -> t -> LSegNode t a
go LSegNode u a
root Int
l0 (Int
l0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) u
forall a. Monoid a => a
mempty)
where
go :: LSegNode t a -> Int -> Int -> t -> LSegNode t a
go LSegNode t a
n Int
l Int
r t
pu | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l Bool -> Bool -> Bool
|| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i = LSegNode t a -> t -> LSegNode t a
forall u a. Action u a => LSegNode u a -> u -> LSegNode u a
applyLSN LSegNode t a
n t
pu
go (LSLeaf a
x) Int
_ Int
_ t
pu = a -> LSegNode t a
forall u a. a -> LSegNode u a
LSLeaf (a -> a
f (a -> t -> a
forall u a. Action u a => a -> u -> a
act a
x t
pu))
go (LSBin a
_ t
u LSegNode t a
lt LSegNode t a
rt) Int
l Int
r t
pu = LSegNode t a -> LSegNode t a -> LSegNode t a
forall u a.
Action u a =>
LSegNode u a -> LSegNode u a -> LSegNode u a
makeLSN (LSegNode t a -> Int -> Int -> t -> LSegNode t a
go LSegNode t a
lt Int
l Int
m t
u') (LSegNode t a -> Int -> Int -> t -> LSegNode t a
go LSegNode t a
rt (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r t
u') where
m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
u' :: t
u' = t
u t -> t -> t
forall a. Semigroup a => a -> a -> a
<> t
pu
updateRangeLST :: Action u a => u -> Int -> Int -> LazySegTree u a -> LazySegTree u a
updateRangeLST :: u -> Int -> Int -> LazySegTree u a -> LazySegTree u a
updateRangeLST u
qu Int
ql Int
qr (LazySegTree lrp :: (Int, Int, Int)
lrp@(Int
l0,Int
r0,Int
p) LSegNode u a
root)
| Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = String -> LazySegTree u a
forall a. HasCallStack => String -> a
error String
"updateRangeLSTM: bad range"
| Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l0 Bool -> Bool -> Bool
|| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
qr = String -> LazySegTree u a
forall a. HasCallStack => String -> a
error String
"updateRangeLSTM: outside range"
| Bool
otherwise = (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
forall u a. (Int, Int, Int) -> LSegNode u a -> LazySegTree u a
LazySegTree (Int, Int, Int)
lrp (LSegNode u a -> Int -> Int -> u -> LSegNode u a
forall a.
Action u a =>
LSegNode u a -> Int -> Int -> u -> LSegNode u a
go LSegNode u a
root Int
l0 (Int
l0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) u
forall a. Monoid a => a
mempty)
where
go :: LSegNode u a -> Int -> Int -> u -> LSegNode u a
go LSegNode u a
n Int
l Int
r u
pu
| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql Bool -> Bool -> Bool
|| Int
qr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = LSegNode u a -> u -> LSegNode u a
forall u a. Action u a => LSegNode u a -> u -> LSegNode u a
applyLSN LSegNode u a
n u
pu
| Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
qr = LSegNode u a -> u -> LSegNode u a
forall u a. Action u a => LSegNode u a -> u -> LSegNode u a
applyLSN LSegNode u a
n (u
pu u -> u -> u
forall a. Semigroup a => a -> a -> a
<> u
qu)
go (LSBin a
_ u
u LSegNode u a
lt LSegNode u a
rt) Int
l Int
r u
pu = LSegNode u a -> LSegNode u a -> LSegNode u a
forall u a.
Action u a =>
LSegNode u a -> LSegNode u a -> LSegNode u a
makeLSN (LSegNode u a -> Int -> Int -> u -> LSegNode u a
go LSegNode u a
lt Int
l Int
m u
u') (LSegNode u a -> Int -> Int -> u -> LSegNode u a
go LSegNode u a
rt (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r u
u') where
m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
u' :: u
u' = u
u u -> u -> u
forall a. Semigroup a => a -> a -> a
<> u
pu
go LSegNode u a
_ Int
_ Int
_ u
_ = String -> LSegNode u a
forall a. HasCallStack => String -> a
error String
"impossible"
foldRangeLST :: Action u a => Int -> Int -> LazySegTree u a -> a
foldRangeLST :: Int -> Int -> LazySegTree u a -> a
foldRangeLST Int
ql Int
qr (LazySegTree (Int
l0,Int
_,Int
p) LSegNode u a
root)
| Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = String -> a
forall a. HasCallStack => String -> a
error String
"foldRangeLST: bad range"
| Bool
otherwise = LSegNode u a -> Int -> Int -> u -> a -> a
forall t a. Action t a => LSegNode t a -> Int -> Int -> t -> a -> a
go LSegNode u a
root Int
l0 (Int
l0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) u
forall a. Monoid a => a
mempty a
forall a. Monoid a => a
mempty
where
go :: LSegNode t a -> Int -> Int -> t -> a -> a
go LSegNode t a
_ Int
l Int
r t
_ a
acc | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql Bool -> Bool -> Bool
|| Int
qr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = a
acc
go (LSLeaf a
x) Int
_ Int
_ t
pu a
acc = a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a -> t -> a
forall u a. Action u a => a -> u -> a
act a
x t
pu
go (LSBin a
x t
u LSegNode t a
lt LSegNode t a
rt) Int
l Int
r t
pu a
acc
| Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
qr = a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a -> t -> a
forall u a. Action u a => a -> u -> a
act a
x t
pu
| Bool
otherwise = LSegNode t a -> Int -> Int -> t -> a -> a
go LSegNode t a
rt (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r t
u' (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! LSegNode t a -> Int -> Int -> t -> a -> a
go LSegNode t a
lt Int
l Int
m t
u' a
acc
where
m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
u' :: t
u' = t
u t -> t -> t
forall a. Semigroup a => a -> a -> a
<> t
pu
foldrLST :: Action u a => (a -> b -> b) -> b -> LazySegTree u a -> b
foldrLST :: (a -> b -> b) -> b -> LazySegTree u a -> b
foldrLST a -> b -> b
f b
z (LazySegTree (Int
l0,Int
r0,Int
p) LSegNode u a
root) = LSegNode u a -> Int -> Int -> u -> b -> b
forall t. Action t a => LSegNode t a -> Int -> Int -> t -> b -> b
go LSegNode u a
root Int
l0 (Int
l0Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) u
forall a. Monoid a => a
mempty b
z where
go :: LSegNode t a -> Int -> Int -> t -> b -> b
go LSegNode t a
_ Int
l Int
_ t
_ | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r0 = b -> b
forall a. a -> a
id
go (LSLeaf a
x) Int
_ Int
_ t
pu = a -> b -> b
f (a -> t -> a
forall u a. Action u a => a -> u -> a
act a
x t
pu)
go (LSBin a
_ t
u LSegNode t a
lt LSegNode t a
rt) Int
l Int
r t
pu = LSegNode t a -> Int -> Int -> t -> b -> b
go LSegNode t a
lt Int
l Int
m t
u' (b -> b) -> (b -> b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LSegNode t a -> Int -> Int -> t -> b -> b
go LSegNode t a
rt (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r t
u' where
m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
u' :: t
u' = t
u t -> t -> t
forall a. Semigroup a => a -> a -> a
<> t
pu
{-# INLINABLE fromListLST #-}
{-# INLINABLE adjustLST #-}
{-# INLINABLE updateRangeLST #-}
{-# INLINABLE foldRangeLST #-}
instance (NFData u, NFData a) => NFData (LazySegTree u a) where
rnf :: LazySegTree u a -> ()
rnf (LazySegTree (Int, Int, Int)
lrp LSegNode u a
n) = (Int, Int, Int) -> ()
forall a. NFData a => a -> ()
rnf (Int, Int, Int)
lrp () -> () -> ()
`seq` LSegNode u a -> ()
forall a. NFData a => a -> ()
rnf LSegNode u a
n
instance (NFData u, NFData a) => NFData (LSegNode u a) where
rnf :: LSegNode u a -> ()
rnf (LSLeaf a
x) = a -> ()
forall a. NFData a => a -> ()
rnf a
x
rnf (LSBin a
x u
u LSegNode u a
lt LSegNode u a
rt) = a -> ()
forall a. NFData a => a -> ()
rnf a
x () -> () -> ()
`seq` u -> ()
forall a. NFData a => a -> ()
rnf u
u () -> () -> ()
`seq` LSegNode u a -> ()
forall a. NFData a => a -> ()
rnf LSegNode u a
lt () -> () -> ()
`seq` LSegNode u a -> ()
forall a. NFData a => a -> ()
rnf LSegNode u a
rt