建树 O(n),修改和查询 O(logn),感觉没有假
是被卡常了吗?
module Main where
type Interval = (Int, Int)
data SegTree = Leaf Int Int | Children Int (SegTree, SegTree) Int Interval deriving Show
-- Leaf val idx | Children val (ls, rs) add_tag (s, t)
getVal :: SegTree -> Int
getVal (Leaf x _) = x
getVal (Children x _ _ _) = x
initTree :: [Int] -> Int -> SegTree
initTree [v] idx = Leaf v idx
initTree l idx = Children (getVal ls + getVal rs) (ls, rs) 0 (idx, idx + length l - 1)
where
mid = length l `div` 2
(s, t) = splitAt mid l
ls = initTree s idx
rs = initTree t (idx + mid)
add :: SegTree -> Interval -> Int -> SegTree
add lf@(Leaf val idx) (l, r) v
| idx >= l && idx <= r = Leaf (val + v) idx
| otherwise = lf
add tr@(Children val (ls, rs) tag (s, t)) (l, r) v
| t < l || s > r = tr
| s >= l && t <= r = Children (val + (t - s + 1) * v) (ls, rs) (tag + v) (s, t)
| otherwise = Children (getVal _ls + getVal _rs + (t - s + 1) * tag) (_ls, _rs) tag (s, t)
where
mid = (s + t) `div` 2
_ls = add ls (l, r) v
_rs = add rs (l, r) v
query :: SegTree -> Interval -> Int
query (Leaf val idx) (l, r)
| idx >=l && idx <= r = val
| otherwise = 0
query (Children val (ls, rs) tag (s, t)) (l, r)
| t < l || s > r = 0
| s >= l && t <= r = val
| otherwise = query ls (l, r) + query rs (l, r) + (min r t - max l s + 1) * tag
getInts :: IO [Int]
getInts = map read . words <$> getLine
solve :: Int -> SegTree -> IO ()
solve 0 _ = return ()
solve m tree = do
op : args <- getInts
if op == 1 then do
let [l, r, k] = args
solve (m - 1) (add tree (l, r) k)
else do
let [l, r] = args
print (query tree (l, r))
solve (m - 1) tree
main :: IO ()
main = do
[n, m] <- getInts
arr <- getInts
let tree = initTree arr 1
solve m tree