{-# LANGUAGE DeriveGeneric, CPP #-}

module Data.KdMap.Dynamic
       ( -- * Usage

         -- $usage

         -- * Reference

         -- ** Types
         PointAsListFn
       , SquaredDistanceFn
       , KdMap
         -- ** Dynamic /k/-d map construction
       , empty
       , singleton
       , emptyWithDist
       , singletonWithDist
         -- ** Insertion
       , insert
       , insertPair
       , batchInsert
         -- ** Query
       , nearest
       , inRadius
       , kNearest
       , inRange
       , assocs
       , keys
       , elems
       , null
       , size
         -- ** Folds
       , foldrWithKey
         -- ** Utilities
       , defaultSqrDist
         -- ** Internal (for testing)
       , subtreeSizes
       ) where

import Prelude hiding (null)

#if MIN_VERSION_base(4,8,0)
#else
import Control.Applicative hiding (empty)
import Data.Foldable
import Data.Traversable
#endif

import Data.Bits
import Data.Function
import qualified Data.List as L

import Control.DeepSeq
import Control.DeepSeq.Generics (genericRnf)
import GHC.Generics

import qualified Data.KdMap.Static as KDM
import Data.KdMap.Static (PointAsListFn, SquaredDistanceFn, defaultSqrDist)

-- $usage
--
-- The 'KdMap' is a variant of
-- @Data.KdTree.Dynamic.@'Data.KdTree.Dynamic.KdTree' where each point
-- in the tree is associated with some data. It is the dynamic variant
-- of @Data.KdMap.Static.@'Data.KdMap.Static.KdMap'.
--
-- Here's an example of interleaving point-value insertions and point
-- queries using 'KdMap', where points are 3D points and values are
-- 'String's:
--
-- @
-- >>> let dkdm = 'singleton' point3dAsList ((Point3D 0.0 0.0 0.0), \"First\")
--
-- >>> let dkdm' = 'insert' dkdm ((Point3D 1.0 1.0 1.0), \"Second\")
--
-- >>> 'nearest' dkdm' (Point3D 0.4 0.4 0.4)
-- (Point3D {x = 0.0, y = 0.0, z = 0.0}, \"First\")
--
-- >>> let dkdm'' = 'insert' dkdm' ((Point3D 0.5 0.5 0.5), \"Third\")
--
-- >>> 'nearest' dkdm'' (Point3D 0.4 0.4 0.4)
-- (Point3D {x = 0.5, y = 0.5, z = 0.5}, \"Third\")
-- @

-- | A dynamic /k/-d tree structure that stores points of type @p@
-- with axis values of type @a@. Additionally, each point is
-- associated with a value of type @v@.
data KdMap a p v = KdMap
                   { KdMap a p v -> [KdMap a p v]
_trees       :: [KDM.KdMap a p v]
                   , KdMap a p v -> PointAsListFn a p
_pointAsList :: PointAsListFn a p
                   , KdMap a p v -> SquaredDistanceFn a p
_distSqr     :: SquaredDistanceFn a p
                   , KdMap a p v -> Int
_numNodes    :: Int
                   } deriving (forall x. KdMap a p v -> Rep (KdMap a p v) x)
-> (forall x. Rep (KdMap a p v) x -> KdMap a p v)
-> Generic (KdMap a p v)
forall x. Rep (KdMap a p v) x -> KdMap a p v
forall x. KdMap a p v -> Rep (KdMap a p v) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a p v x. Rep (KdMap a p v) x -> KdMap a p v
forall a p v x. KdMap a p v -> Rep (KdMap a p v) x
$cto :: forall a p v x. Rep (KdMap a p v) x -> KdMap a p v
$cfrom :: forall a p v x. KdMap a p v -> Rep (KdMap a p v) x
Generic
instance (NFData a, NFData p, NFData v) => NFData (KdMap a p v) where rnf :: KdMap a p v -> ()
rnf = KdMap a p v -> ()
forall a. (Generic a, GNFData (Rep a)) => a -> ()
genericRnf

instance (Show a, Show p, Show v) => Show (KdMap a p v) where
  show :: KdMap a p v -> String
show KdMap a p v
kdm = String
"KdMap " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [KdMap a p v] -> String
forall a. Show a => a -> String
show (KdMap a p v -> [KdMap a p v]
forall a p v. KdMap a p v -> [KdMap a p v]
_trees KdMap a p v
kdm)

instance Functor (KdMap a p) where
  fmap :: (a -> b) -> KdMap a p a -> KdMap a p b
fmap a -> b
f KdMap a p a
dkdMap = KdMap a p a
dkdMap { _trees :: [KdMap a p b]
_trees = (KdMap a p a -> KdMap a p b) -> [KdMap a p a] -> [KdMap a p b]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> KdMap a p a -> KdMap a p b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) ([KdMap a p a] -> [KdMap a p b]) -> [KdMap a p a] -> [KdMap a p b]
forall a b. (a -> b) -> a -> b
$ KdMap a p a -> [KdMap a p a]
forall a p v. KdMap a p v -> [KdMap a p v]
_trees KdMap a p a
dkdMap }

-- | Performs a foldr over each point-value pair in the 'KdMap'.
foldrWithKey :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey (p, v) -> b -> b
f b
z KdMap a p v
dkdMap = (KdMap a p v -> b -> b) -> b -> [KdMap a p v] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr ((b -> KdMap a p v -> b) -> KdMap a p v -> b -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((b -> KdMap a p v -> b) -> KdMap a p v -> b -> b)
-> (b -> KdMap a p v -> b) -> KdMap a p v -> b -> b
forall a b. (a -> b) -> a -> b
$ ((p, v) -> b -> b) -> b -> KdMap a p v -> b
forall p v b a. ((p, v) -> b -> b) -> b -> KdMap a p v -> b
KDM.foldrWithKey (p, v) -> b -> b
f) b
z ([KdMap a p v] -> b) -> [KdMap a p v] -> b
forall a b. (a -> b) -> a -> b
$ KdMap a p v -> [KdMap a p v]
forall a p v. KdMap a p v -> [KdMap a p v]
_trees KdMap a p v
dkdMap

instance Foldable (KdMap a p) where
  foldr :: (a -> b -> b) -> b -> KdMap a p a -> b
foldr a -> b -> b
f = ((p, a) -> b -> b) -> b -> KdMap a p a -> b
forall p v b a. ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey (a -> b -> b
f (a -> b -> b) -> ((p, a) -> a) -> (p, a) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, a) -> a
forall a b. (a, b) -> b
snd)

instance Traversable (KdMap a p) where
  traverse :: (a -> f b) -> KdMap a p a -> f (KdMap a p b)
traverse a -> f b
f (KdMap [KdMap a p a]
t PointAsListFn a p
p SquaredDistanceFn a p
d Int
n) =
    [KdMap a p b]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p b
forall a p v.
[KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
KdMap ([KdMap a p b]
 -> PointAsListFn a p
 -> SquaredDistanceFn a p
 -> Int
 -> KdMap a p b)
-> f [KdMap a p b]
-> f (PointAsListFn a p
      -> SquaredDistanceFn a p -> Int -> KdMap a p b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (KdMap a p a -> f (KdMap a p b))
-> [KdMap a p a] -> f [KdMap a p b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((a -> f b) -> KdMap a p a -> f (KdMap a p b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f) [KdMap a p a]
t f (PointAsListFn a p
   -> SquaredDistanceFn a p -> Int -> KdMap a p b)
-> f (PointAsListFn a p)
-> f (SquaredDistanceFn a p -> Int -> KdMap a p b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PointAsListFn a p -> f (PointAsListFn a p)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PointAsListFn a p
p f (SquaredDistanceFn a p -> Int -> KdMap a p b)
-> f (SquaredDistanceFn a p) -> f (Int -> KdMap a p b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SquaredDistanceFn a p -> f (SquaredDistanceFn a p)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SquaredDistanceFn a p
d f (Int -> KdMap a p b) -> f Int -> f (KdMap a p b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> f Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n

-- | Generates an empty 'KdMap' with a user-specified distance function.
emptyWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 = [KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
forall a p v.
[KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
KdMap [] PointAsListFn a p
p2l SquaredDistanceFn a p
d2 Int
0

-- | Returns whether the 'KdMap' is empty.
null :: KdMap a p v -> Bool
null :: KdMap a p v -> Bool
null (KdMap [] PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
_) = Bool
True
null KdMap a p v
_ = Bool
False

-- | Generates a 'KdMap' with a single point-value pair using a
-- user-specified distance function.
singletonWithDist :: Real a => PointAsListFn a p
                            -> SquaredDistanceFn a p
                            -> (p, v)
                            -> KdMap a p v
singletonWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> (p, v) -> KdMap a p v
singletonWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 (p
k, v
v) =
  [KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
forall a p v.
[KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
KdMap [PointAsListFn a p
-> SquaredDistanceFn a p -> [(p, v)] -> KdMap a p v
forall a p v.
Real a =>
PointAsListFn a p
-> SquaredDistanceFn a p -> [(p, v)] -> KdMap a p v
KDM.buildWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 [(p
k, v
v)]] PointAsListFn a p
p2l SquaredDistanceFn a p
d2 Int
1

-- | Generates an empty 'KdMap' with the default distance function.
empty :: Real a => PointAsListFn a p -> KdMap a p v
empty :: PointAsListFn a p -> KdMap a p v
empty PointAsListFn a p
p2l = PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
forall a p v.
PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist PointAsListFn a p
p2l (SquaredDistanceFn a p -> KdMap a p v)
-> SquaredDistanceFn a p -> KdMap a p v
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l

-- | Generates a 'KdMap' with a single point-value pair using the
-- default distance function.
singleton :: Real a => PointAsListFn a p -> (p, v) -> KdMap a p v
singleton :: PointAsListFn a p -> (p, v) -> KdMap a p v
singleton PointAsListFn a p
p2l = PointAsListFn a p -> SquaredDistanceFn a p -> (p, v) -> KdMap a p v
forall a p v.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> (p, v) -> KdMap a p v
singletonWithDist PointAsListFn a p
p2l (SquaredDistanceFn a p -> (p, v) -> KdMap a p v)
-> SquaredDistanceFn a p -> (p, v) -> KdMap a p v
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l

-- | Adds a given point-value pair to a 'KdMap'.
--
-- Average time complexity per insert for /n/ inserts: /O(log^2(n))/.
insert :: Real a => KdMap a p v -> p -> v -> KdMap a p v
insert :: KdMap a p v -> p -> v -> KdMap a p v
insert (KdMap [KdMap a p v]
trees PointAsListFn a p
p2l SquaredDistanceFn a p
d2 Int
n) p
k v
v =
  let bitList :: [Int]
bitList = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int
1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&.) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR`)) [Int
0..]
      ([(Int, KdMap a p v)]
onesPairs, [(Int, KdMap a p v)]
theRestPairs) = ((Int, KdMap a p v) -> Bool)
-> [(Int, KdMap a p v)]
-> ([(Int, KdMap a p v)], [(Int, KdMap a p v)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) (Int -> Bool)
-> ((Int, KdMap a p v) -> Int) -> (Int, KdMap a p v) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, KdMap a p v) -> Int
forall a b. (a, b) -> a
fst) ([(Int, KdMap a p v)]
 -> ([(Int, KdMap a p v)], [(Int, KdMap a p v)]))
-> [(Int, KdMap a p v)]
-> ([(Int, KdMap a p v)], [(Int, KdMap a p v)])
forall a b. (a -> b) -> a -> b
$ [Int] -> [KdMap a p v] -> [(Int, KdMap a p v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
bitList [KdMap a p v]
trees
      (([Int]
_, [KdMap a p v]
ones), ([Int]
_, [KdMap a p v]
theRest)) = ([(Int, KdMap a p v)] -> ([Int], [KdMap a p v])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Int, KdMap a p v)]
onesPairs, [(Int, KdMap a p v)] -> ([Int], [KdMap a p v])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Int, KdMap a p v)]
theRestPairs)
      newTree :: KdMap a p v
newTree = PointAsListFn a p
-> SquaredDistanceFn a p -> [(p, v)] -> KdMap a p v
forall a p v.
Real a =>
PointAsListFn a p
-> SquaredDistanceFn a p -> [(p, v)] -> KdMap a p v
KDM.buildWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2  ([(p, v)] -> KdMap a p v) -> [(p, v)] -> KdMap a p v
forall a b. (a -> b) -> a -> b
$ (p
k, v
v) (p, v) -> [(p, v)] -> [(p, v)]
forall a. a -> [a] -> [a]
: (KdMap a p v -> [(p, v)]) -> [KdMap a p v] -> [(p, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
L.concatMap KdMap a p v -> [(p, v)]
forall a p v. KdMap a p v -> [(p, v)]
KDM.assocs [KdMap a p v]
ones
  in  [KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
forall a p v.
[KdMap a p v]
-> PointAsListFn a p -> SquaredDistanceFn a p -> Int -> KdMap a p v
KdMap (KdMap a p v
newTree KdMap a p v -> [KdMap a p v] -> [KdMap a p v]
forall a. a -> [a] -> [a]
: [KdMap a p v]
theRest) PointAsListFn a p
p2l SquaredDistanceFn a p
d2 (Int -> KdMap a p v) -> Int -> KdMap a p v
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- | Same as 'insert', but takes point and value as a pair.
insertPair :: Real a => KdMap a p v -> (p, v) -> KdMap a p v
insertPair :: KdMap a p v -> (p, v) -> KdMap a p v
insertPair KdMap a p v
t = (p -> v -> KdMap a p v) -> (p, v) -> KdMap a p v
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (KdMap a p v -> p -> v -> KdMap a p v
forall a p v. Real a => KdMap a p v -> p -> v -> KdMap a p v
insert KdMap a p v
t)

-- | Given a 'KdMap' and a query point, returns the point-value pair in
-- the 'KdMap' with the point nearest to the query.
--
-- Average time complexity: /O(log^2(n))/.
nearest :: Real a => KdMap a p v -> p -> (p, v)
nearest :: KdMap a p v -> p -> (p, v)
nearest (KdMap [KdMap a p v]
ts PointAsListFn a p
_ SquaredDistanceFn a p
d2 Int
_) p
query =
  let nearests :: [(p, v)]
nearests = (KdMap a p v -> (p, v)) -> [KdMap a p v] -> [(p, v)]
forall a b. (a -> b) -> [a] -> [b]
map (KdMap a p v -> p -> (p, v)
forall a p v. Real a => KdMap a p v -> p -> (p, v)
`KDM.nearest` p
query) [KdMap a p v]
ts
  --in  if   Data.List.null nearests
  in  if [(p, v)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [(p, v)]
nearests
      then String -> (p, v)
forall a. HasCallStack => String -> a
error String
"Called nearest on empty KdMap."
      else ((p, v) -> (p, v) -> Ordering) -> [(p, v)] -> (p, v)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
L.minimumBy (a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (a -> a -> Ordering)
-> ((p, v) -> a) -> (p, v) -> (p, v) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (SquaredDistanceFn a p
d2 p
query (p -> a) -> ((p, v) -> p) -> (p, v) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, v) -> p
forall a b. (a, b) -> a
fst)) [(p, v)]
nearests

-- | Given a 'KdMap', a query point, and a number @k@, returns the
-- @k@ point-value pairs with the nearest points to the query.
--
-- Neighbors are returned in order of increasing distance from query
-- point.
--
-- Average time complexity: /log(k) * log^2(n)/ for /k/ nearest
-- neighbors on a structure with /n/ data points.
--
-- Worst case time complexity: /n * log(k)/ for /k/ nearest neighbors
-- on a structure with /n/ data points.
kNearest :: Real a => KdMap a p v -> Int -> p -> [(p, v)]
kNearest :: KdMap a p v -> Int -> p -> [(p, v)]
kNearest (KdMap [KdMap a p v]
trees PointAsListFn a p
_ SquaredDistanceFn a p
d2 Int
_) Int
k p
query =
  let neighborSets :: [[(p, v)]]
neighborSets = (KdMap a p v -> [(p, v)]) -> [KdMap a p v] -> [[(p, v)]]
forall a b. (a -> b) -> [a] -> [b]
map (\KdMap a p v
t -> KdMap a p v -> Int -> p -> [(p, v)]
forall a p v. Real a => KdMap a p v -> Int -> p -> [(p, v)]
KDM.kNearest KdMap a p v
t Int
k p
query) [KdMap a p v]
trees
  in  Int -> [(p, v)] -> [(p, v)]
forall a. Int -> [a] -> [a]
take Int
k ([(p, v)] -> [(p, v)]) -> [(p, v)] -> [(p, v)]
forall a b. (a -> b) -> a -> b
$ ([(p, v)] -> [(p, v)] -> [(p, v)])
-> [(p, v)] -> [[(p, v)]] -> [(p, v)]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr [(p, v)] -> [(p, v)] -> [(p, v)]
forall b. [(p, b)] -> [(p, b)] -> [(p, b)]
merge [] [[(p, v)]]
neighborSets
 where merge :: [(p, b)] -> [(p, b)] -> [(p, b)]
merge [] [(p, b)]
ys = [(p, b)]
ys
       merge [(p, b)]
xs [] = [(p, b)]
xs
       merge xs :: [(p, b)]
xs@((p, b)
x:[(p, b)]
xt) ys :: [(p, b)]
ys@((p, b)
y:[(p, b)]
yt)
         | a
distX a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
distY = (p, b)
x (p, b) -> [(p, b)] -> [(p, b)]
forall a. a -> [a] -> [a]
: [(p, b)] -> [(p, b)] -> [(p, b)]
merge [(p, b)]
xt [(p, b)]
ys
         | Bool
otherwise      = (p, b)
y (p, b) -> [(p, b)] -> [(p, b)]
forall a. a -> [a] -> [a]
: [(p, b)] -> [(p, b)] -> [(p, b)]
merge [(p, b)]
xs [(p, b)]
yt
        where distX :: a
distX = SquaredDistanceFn a p
d2 p
query (p -> a) -> p -> a
forall a b. (a -> b) -> a -> b
$ (p, b) -> p
forall a b. (a, b) -> a
fst (p, b)
x
              distY :: a
distY = SquaredDistanceFn a p
d2 p
query (p -> a) -> p -> a
forall a b. (a -> b) -> a -> b
$ (p, b) -> p
forall a b. (a, b) -> a
fst (p, b)
y

-- | Given a 'KdMap', a query point, and a radius, returns all
-- point-value pairs in the 'KdTree' with points within the given
-- radius of the query point.
--
-- Points are not returned in any particular order.
--
-- Worst case time complexity: /O(n)/ for /n/ data points.
inRadius :: Real a => KdMap a p v -> a -> p -> [(p, v)]
inRadius :: KdMap a p v -> a -> p -> [(p, v)]
inRadius (KdMap [KdMap a p v]
trees PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
_) a
radius p
query =
  (KdMap a p v -> [(p, v)]) -> [KdMap a p v] -> [(p, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
L.concatMap (\KdMap a p v
t -> KdMap a p v -> a -> p -> [(p, v)]
forall a p v. Real a => KdMap a p v -> a -> p -> [(p, v)]
KDM.inRadius KdMap a p v
t a
radius p
query) [KdMap a p v]
trees

-- | Finds all point-value pairs in a 'KdMap' with points within a
-- given range, where the range is specified as a set of lower and
-- upper bounds.
--
-- Points are not returned in any particular order.
--
-- Worst case time complexity: /O(n)/ for n data points and a range
-- that spans all the points.
inRange :: Real a => KdMap a p v
                  -> p -- ^ lower bounds of range
                  -> p -- ^ upper bounds of range
                  -> [(p, v)] -- ^ point-value pairs within given
                              -- range
inRange :: KdMap a p v -> p -> p -> [(p, v)]
inRange (KdMap [KdMap a p v]
trees PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
_) p
lowers p
uppers =
  (KdMap a p v -> [(p, v)]) -> [KdMap a p v] -> [(p, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
L.concatMap (\KdMap a p v
t -> KdMap a p v -> p -> p -> [(p, v)]
forall a p v. Real a => KdMap a p v -> p -> p -> [(p, v)]
KDM.inRange KdMap a p v
t p
lowers p
uppers) [KdMap a p v]
trees

-- | Returns the number of elements in the 'KdMap'.
--
-- Time complexity: /O(1)/
size :: KdMap a p v -> Int
size :: KdMap a p v -> Int
size (KdMap [KdMap a p v]
_ PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
n) = Int
n

-- | Returns a list of all the point-value pairs in the 'KdMap'.
--
-- Time complexity: /O(n)/ for /n/ data points.
assocs :: KdMap a p v -> [(p, v)]
assocs :: KdMap a p v -> [(p, v)]
assocs (KdMap [KdMap a p v]
trees PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
_) = (KdMap a p v -> [(p, v)]) -> [KdMap a p v] -> [(p, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
L.concatMap KdMap a p v -> [(p, v)]
forall a p v. KdMap a p v -> [(p, v)]
KDM.assocs [KdMap a p v]
trees

-- | Returns all points in the 'KdMap'.
--
-- Time complexity: /O(n)/ for /n/ data points.
keys :: KdMap a p v -> [p]
keys :: KdMap a p v -> [p]
keys = ((p, v) -> p) -> [(p, v)] -> [p]
forall a b. (a -> b) -> [a] -> [b]
map (p, v) -> p
forall a b. (a, b) -> a
fst ([(p, v)] -> [p])
-> (KdMap a p v -> [(p, v)]) -> KdMap a p v -> [p]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KdMap a p v -> [(p, v)]
forall a p v. KdMap a p v -> [(p, v)]
assocs

-- | Returns all values in the 'KdMap'.
--
-- Time complexity: /O(n)/ for /n/ data points.
elems :: KdMap a p v -> [v]
elems :: KdMap a p v -> [v]
elems = ((p, v) -> v) -> [(p, v)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (p, v) -> v
forall a b. (a, b) -> b
snd ([(p, v)] -> [v])
-> (KdMap a p v -> [(p, v)]) -> KdMap a p v -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KdMap a p v -> [(p, v)]
forall a p v. KdMap a p v -> [(p, v)]
assocs

-- | Inserts a list of point-value pairs into the 'KdMap'.
--
-- TODO: This will be made far more efficient than simply repeatedly
-- inserting.
batchInsert :: Real a => KdMap a p v -> [(p, v)] -> KdMap a p v
batchInsert :: KdMap a p v -> [(p, v)] -> KdMap a p v
batchInsert =  (KdMap a p v -> (p, v) -> KdMap a p v)
-> KdMap a p v -> [(p, v)] -> KdMap a p v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' KdMap a p v -> (p, v) -> KdMap a p v
forall a p v. Real a => KdMap a p v -> (p, v) -> KdMap a p v
insertPair

-- | Returns size of each internal /k/-d tree that makes up the
-- dynamic structure. For internal testing use.
subtreeSizes :: KdMap a p v -> [Int]
subtreeSizes :: KdMap a p v -> [Int]
subtreeSizes (KdMap [KdMap a p v]
trees PointAsListFn a p
_ SquaredDistanceFn a p
_ Int
_) = (KdMap a p v -> Int) -> [KdMap a p v] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map KdMap a p v -> Int
forall a p v. KdMap a p v -> Int
KDM.size [KdMap a p v]
trees