segment tree build and query

segment_tree.hs

-- video https://www.youtube.com/watch?v=nckWiHmeNXU&feature=youtu.be
{-# LANGUAGE RankNTypes #-}
-- GHCi, version 8.8.4
-- data Btree a = Empty | Node a (Btree a) (Btree a) deriving Show
-- https://www.youtube.com/watch?v=Ud-1Z0hBlB8&t=379s

-- data Btree a = Empty | Node a (Btree a) (Btree a) deriving Show
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
import PrettyT -- https://www.youtube.com/watch?v=Ud-1Z0hBlB8&t=3s
import Data.List.Split (chunksOf, whenElt)
import Data.Semigroup (Sum (Sum), Product (Product), Min (Min), Max (Max))

-- Segment Tree
-- binary tree 
-- leaves correspond to the elementary intervals (points)
-- The internal nodes of T correspond to intervals that are the union of elementary intervals
-- Each node in T stores the segments that span through its interval, 
-- but do not span through the interval of its parent (only part).

-- build perfect tree (but no nodes needed if where is no data) and its height

build :: Monoid a => [a] -> (Int,Btree a)
build l = go 0 l'' where
  l' = if even $ length l then l else l++[mempty]
  l'' = [Node e Empty Empty | e <- l']  
  go h [r] = (h,r)
  go h l = go (h+1) el
    where
      l_ = [b c| c <- chunksOf 2 l]
      el | (even $ length l_) || (length l_==1) = l_ 
         | otherwise  = l_++[Node mempty Empty Empty]
         
      

  b (n1@(Node e1 _ _):n2@(Node e2 _ _):[]) = Node (mappend e1 e2) n1 n2
  b (n1:[]) = n1
  b [] = error "empty chunks"
  
query :: forall a.Monoid a => Int -> Btree a -> Int -> Int -> a
-- l - number of leafs in our perfect tree
query l t a b | not ((a>=0) && (b>=a) && (b<=l)) = error "invalid interval"
              | otherwise = go t l 0 (l-1)
              where
                go :: Monoid a => Btree a -> Int -> Int -> Int -> a 
                go t@(Node e l r) ln ta tb
                  | inside = e
                  | outside = mempty
                  | otherwise  = mappend l' r' -- partial
                  where
                    -- perfect tree expected here
                    l' = go l ld2 ta (ta+((ln `div` 2)-1))
                    r' = go r ld2 (ta+(ln `div` 2)) tb
                    ld2 = ln `div` 2
                    inside = (ta >= a) && (tb <= b)
                    outside = (a>tb) || (b<ta)
                    
                    -- partial =         
--l = [Sum 1,Sum 2, Sum 3, Sum 4]

q ti a b = query (2^(fst ti)) (snd ti) a b


--l = [Product 1,Product 2, Product 3, Product 4]
-- l :: [Min Int]
-- l = [Min 1, Min 2, Min 3, Min 4]
l :: [Max Int]
l = [Max 1, Max 2, Max 3, Max 4]
t = build l
-- build O(N)
-- aux space O(N) (N on bottom + N/2+N/4+N/8.. = N+N(1/2+1/4+1/8..) = N+N*1=2N)
-- update Log(N)
-- request Log(N)
-- https://www.youtube.com/watch?v=wyt5_GJqsMw
-- https://en.wikipedia.org/wiki/Range_query_(data_structures)
-- https://en.wikipedia.org/wiki/Segment_tree

https://www.youtube.com/watch?v=nckWiHmeNXU&feature=youtu.be