hgeometry-0.9.0.0: src/Data/Geometry/SegmentTree/Generic.hs
{-# LANGUAGE TemplateHaskell #-}
--------------------------------------------------------------------------------
-- |
-- Module : Data.Geometry.SegmentTree.Generic
-- Copyright : (C) Frank Staals
-- License : see the LICENSE file
-- Maintainer : Frank Staals
-- Description : Implementation of SegmentTrees
--
--------------------------------------------------------------------------------
module Data.Geometry.SegmentTree.Generic( NodeData(..), splitPoint, range, assoc
, LeafData(..), atomicRange, leafAssoc
, SegmentTree(..), unSegmentTree
, Assoc(..)
, createTree, fromIntervals
, insert, delete
, search, stab
, I(..), fromIntervals'
, Count(..)
) where
import Control.DeepSeq
import Control.Lens
import Data.BinaryTree
import Data.Geometry.Interval
import Data.Geometry.IntervalTree (IntervalLike(..))
import Data.Geometry.Properties
import qualified Data.List as List
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import GHC.Generics (Generic)
--------------------------------------------------------------------------------
-- | Internal nodes store a split point, the range, and an associated data structure
data NodeData v r = NodeData { _splitPoint :: !(EndPoint r)
, _range :: !(Range r)
, _assoc :: !v
} deriving (Show,Eq,Functor,Generic)
makeLenses ''NodeData
instance (NFData v, NFData r) => NFData (NodeData v r)
-- | We store atomic ranges a bit more efficiently.
data AtomicRange r = Singleton !r | AtomicRange deriving (Show,Eq,Functor,Generic)
instance NFData r => NFData (AtomicRange r)
-- | Leaf nodes store an atomic range, and an associated data structure.
data LeafData v r = LeafData { _atomicRange :: !(AtomicRange r)
, _leafAssoc :: !v
} deriving (Show,Eq,Functor,Generic)
makeLenses ''LeafData
instance (NFData v, NFData r) => NFData (LeafData v r)
--------------------------------------------------------------------------------
-- | Segment tree on a Fixed set of endpoints
newtype SegmentTree v r =
SegmentTree { _unSegmentTree :: BinLeafTree (NodeData v r) (LeafData v r) }
deriving (Show,Eq,Generic,NFData)
makeLenses ''SegmentTree
-- rangeOf :: BinLeafTree (NodeData v r) (LeafData v r) -> Range (UnBounded r)
-- rangeOf (Node _ x _) = Val <$> x^.range
-- rangeOf (Leaf x) = case x^.atomicRange of
-- Singleton r -> ClsoedRange (Val r) (Val r)
-- AtomicRange -> OpenRange MinInfinity MaxInfinity
data BuildLeaf a = LeafSingleton a | LeafRange a a deriving (Show,Eq)
-- | Given a sorted list of endpoints, without duplicates, construct a segment tree
--
--
-- \(O(n)\) time
createTree :: NonEmpty r -> v -> SegmentTree v r
-- createTree (r NonEmpty.:| []) v = SegmentTree . Leaf $ LeafData (Singleton r) v
createTree pts v = SegmentTree . fmap h . foldUpData f g . fmap _unElem
. asBalancedBinLeafTree $ ranges
where
h (LeafSingleton r) = LeafData (Singleton r) v
h (LeafRange _ _) = LeafData AtomicRange v
f l _ r = let m = l^.range.upper
ll = l^.range.lower
rr = r^.range.upper
in NodeData m (Range ll rr) v
-- | Singletons map to closed singleton ranges, Ranges map to open ranges
g (LeafSingleton r) = NodeData (Closed r) (ClosedRange r r) v
g (LeafRange s r) = NodeData (Open r) (OpenRange s r) v
ranges = interleave (fmap LeafSingleton pts) ranges'
ranges' = zipWith LeafRange (NonEmpty.toList pts) (NonEmpty.tail pts)
-- | Interleaves the two lists
--
-- >>> interleave (NonEmpty.fromList ["0","1","2"]) ["01","12"]
-- "0" :| ["01","1","12","2"]
interleave :: NonEmpty a -> [a] -> NonEmpty a
interleave (x NonEmpty.:| xs) ys = x NonEmpty.:| concat (zipWith (\a b -> [a,b]) ys xs)
-- | Build a SegmentTree
--
-- \(O(n \log n)\)
fromIntervals :: (Ord r, Eq p, Assoc v i, IntervalLike i, Monoid v, NumType i ~ r)
=> (Interval p r -> i)
-> NonEmpty (Interval p r) -> SegmentTree v r
fromIntervals f is = foldr (insert . f) (createTree pts mempty) is
where
endPoints (toRange -> Range' a b) = [a,b]
pts = nub' . NonEmpty.sort . NonEmpty.fromList . concatMap endPoints $ is
nub' = fmap NonEmpty.head . NonEmpty.group1
-- -- | lists all intervals
-- toList :: SegmentTree v r -> [i]
-- toList = undefined
--------------------------------------------------------------------------------
-- * Searching
-- | Search for all intervals intersecting x
--
-- \(O(\log n + k)\) where \(k\) is the output size
search :: (Ord r, Monoid v) => r -> SegmentTree v r -> v
search x = mconcat . stab x
inAtomicRange :: Eq r => r -> AtomicRange r -> Bool
x `inAtomicRange` (Singleton r) = x == r
_ `inAtomicRange` AtomicRange = True
-- | Returns the associated values of the nodes on the search path to x
--
-- \(O(\log n)\)
stab :: Ord r => r -> SegmentTree v r -> [v]
stab x (SegmentTree t) = stabRoot t
where
stabRoot (Leaf (LeafData rr v))
| x `inAtomicRange` rr = [v]
| otherwise = []
stabRoot (Node l (NodeData m rr v) r) = case (x `inRange` rr, Closed x <= m) of
(False,_) -> []
(True,True) -> v : stab' l
_ -> v : stab' r
stab' (Leaf (LeafData rr v)) | x `inAtomicRange` rr = [v]
| otherwise = []
stab' (Node l (NodeData m _ v) r) | Closed x <= m = v : stab' l
| otherwise = v : stab' r
--------------------------------------------------------------------------------
-- * Inserting intervals
-- | Class for associcated data structures
class Measured v i => Assoc v i where
insertAssoc :: i -> v -> v
deleteAssoc :: i -> v -> v
-- | Gets the range associated with this node
getRange :: BinLeafTree (NodeData v r) (LeafData t r) -> Maybe (Range r)
getRange (Leaf (LeafData (Singleton r) _)) = Just $ Range (Closed r) (Closed r)
getRange (Leaf _) = Nothing
getRange (Node _ nd _) = Just $ nd^.range
coversAtomic :: Ord r
=> Range r -> Range r -> AtomicRange r -> Bool
coversAtomic ri _ (Singleton r) = r `inRange` ri
coversAtomic ri inR AtomicRange = ri `covers` inR
-- | Pre: the interval should have one of the endpoints on which the tree is built.
insert :: (Assoc v i, NumType i ~ r, Ord r, IntervalLike i)
=> i -> SegmentTree v r -> SegmentTree v r
insert i (SegmentTree t) = SegmentTree $ insertRoot t
where
ri@(Range a b) = toRange i
insertRoot t' = maybe t' (flip insert' t') $ getRange t'
insert' inR lf@(Leaf nd@(LeafData rr _))
| coversAtomic ri inR rr = Leaf $ nd&leafAssoc %~ insertAssoc i
| otherwise = lf
insert' (Range c d) (Node l nd@(NodeData m rr _) r)
| ri `covers` rr = Node l (nd&assoc %~ insertAssoc i) r
| otherwise = Node l' nd r'
where
-- check if the range intersects the range of the left subtree
l' = if a <= m then insert' (Range c m) l else l
r' = if m < b then insert' (Range (toOpen m) d) r else r
toOpen = Open . view unEndPoint
-- | Delete an interval from the tree
--
-- pre: The segment is in the tree!
delete :: (Assoc v i, NumType i ~ r, Ord r, IntervalLike i)
=> i -> SegmentTree v r -> SegmentTree v r
delete i (SegmentTree t) = SegmentTree $ delete' t
where
(Range _ b) = toRange i
delete' (Leaf ld) = Leaf $ ld&leafAssoc %~ deleteAssoc i
delete' (Node l nd@(_splitPoint -> m) r)
| b <= m = Node (delete' l) (nd&assoc %~ deleteAssoc i) r
| otherwise = Node l (nd&assoc %~ deleteAssoc i) (delete' r)
-- delete'' (Leaf ld) = Leaf $ ld&leafAssoc %~ deleteAssoc i
-- delete'' (Node l nd r) = Node l (nd&assoc %~ deleteAssoc i) r
-- deleteL (Leaf ld) = Leaf $ ld&leafAssoc %~ deleteAssoc i
-- deleteL (Node l nd@(_splitPoint -> m) r)
-- | a <= m = Node (deleteL l) (nd&assoc %~ deleteAssoc i) (delete'' r)
-- | otherwise = Node l nd (deleteL r)
-- deleteR (Leaf ld) = Leaf $ ld&leafAssoc %~ deleteAssoc i
-- deleteR (Node l nd@(_splitPoint -> m) r)
-- | m <= b = Node (delete'' l) (nd&assoc %~ deleteAssoc i) (deleteR r)
-- | otherwise = Node (deleteR l) nd r
--------------------------------------------------------------------------------
-- * Listing the intervals stabbed
-- | Interval
newtype I a = I { _unI :: a} deriving (Show,Read,Eq,Ord,Generic,NFData)
type instance NumType (I a) = NumType a
instance Measured [I a] (I a) where
measure = (:[])
instance Eq a => Assoc [I a] (I a) where
insertAssoc = (:)
deleteAssoc = List.delete
-- instance Measured [Interval p r] (Interval p r) where
-- measure = (:[])
-- instance (Eq p, Eq r) => Assoc [Interval p r] (Interval p r) where
-- insertAssoc = (:)
-- deleteAssoc = List.delete
instance IntervalLike a => IntervalLike (I a) where
toRange = toRange . _unI
fromIntervals' :: (Eq p, Ord r)
=> NonEmpty (Interval p r) -> SegmentTree [I (Interval p r)] r
fromIntervals' = fromIntervals I
--------------------------------------------------------------------------------
-- * Counting the number of segments intersected
newtype Count = Count { getCount :: Word }
deriving (Show,Eq,Ord,Num,Integral,Enum,Real,Generic,NFData)
newtype C a = C { _unC :: a} deriving (Show,Read,Eq,Ord,Generic,NFData)
instance Semigroup Count where
a <> b = Count $ getCount a + getCount b
instance Monoid Count where
mempty = 0
mappend = (<>)
instance Measured Count (C i) where
measure _ = 1
instance Assoc Count (C i) where
insertAssoc _ v = v + 1
deleteAssoc _ v = v - 1
--------------------------------------------------------------------------------
-- * Testing stuff
-- test'' = fromIntervals' . NonEmpty.fromList $ test
-- test = [Interval (Closed (238 :+ ())) (Open (309 :+ ())), Interval (Closed (175 :+ ())) (Closed (269 :+ ())),Interval (Closed (255 :+ ())) (Open (867 :+ ())),Interval (Open (236 :+ ())) (Closed (863 :+ ())),Interval (Open (150 :+ ())) (Closed (161 :+ ())),Interval (Closed (35 :+ ())) (Closed (77 :+ ()))]
-- -- q = [78]
-- -- test = fromIntervals' . NonEmpty.fromList $ [ closedInterval 0 10
-- -- , closedInterval 5 15
-- -- , closedInterval 1 4
-- -- , closedInterval 3 9
-- -- ]
-- tst = fromIntervals' . NonEmpty.fromList $ [ closedInterval 1 6
-- , closedInterval 2 6
-- -- , Interval (Closed $ ext 0) (Open $ ext 1)
-- ]
-- closedInterval a b = ClosedInterval (ext a) (ext b)
-- showT :: (Show r, Show v) => SegmentTree v r -> String
-- showT = drawTree . _unSegmentTree
-- test' :: (Show r, Num r, Ord r, Enum r) => SegmentTree [I (Interval () r)] r
-- test' = insert (I $ closedInterval 6 14) $ createTree (NonEmpty.fromList [2,4..20]) []