hgeometry-0.14: 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 Data.Measured.Class
import Data.Measured.Size
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 (asRange -> 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) = asRange i
insertRoot t' = maybe t' (`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) = asRange 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
asRange = asRange . _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]) []