{-# LANGUAGE DeriveGeneric, CPP, FlexibleContexts #-}

module Data.KdMap.Static
       ( -- * Usage

         -- $usage

         -- * Reference

         -- ** Types
         PointAsListFn
       , SquaredDistanceFn
       , KdMap
         -- ** /k/-d map construction
       , empty
       , emptyWithDist
       , singleton
       , singletonWithDist
       , build
       , buildWithDist
       , insertUnbalanced
       , batchInsertUnbalanced
         -- ** Query
       , nearest
       , inRadius
       , kNearest
       , inRange
       , assocs
       , keys
       , elems
       , null
       , size
         -- ** Folds
       , foldrWithKey
         -- ** Utilities
       , defaultSqrDist
         -- ** Advanced
       , TreeNode(..)
         -- ** Internal (for testing)
       , isValid
       ) where

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

import Control.Applicative hiding (empty)

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

import Prelude hiding (null)
import qualified Data.List as L
import Data.Maybe
import Data.Ord
import qualified Data.Heap as Q

-- $usage
--
-- The 'KdMap' is a variant of 'Data.KdTree.Static.KdTree' where each point in
-- the tree is associated with some data. When talking about 'KdMap's,
-- we'll refer to the points and their associated data as the /points/
-- and /values/ of the 'KdMap', respectively. It might help to think
-- of 'Data.KdTree.Static.KdTree' and 'KdMap' as being analogous to
-- 'Set' and 'Map'.
--
-- Suppose you wanted to perform point queries on a set of 3D points,
-- where each point is associated with a 'String'. Here's how to build
-- a 'KdMap' of the data and perform a nearest neighbor query (if this
-- doesn't make sense, start with the documentation for
-- 'Data.KdTree.Static.KdTree'):
--
-- @
-- >>> let points = [(Point3d 0.0 0.0 0.0), (Point3d 1.0 1.0 1.0)]
--
-- >>> let valueStrings = [\"First\", \"Second\"]
--
-- >>> let pointValuePairs = 'zip' points valueStrings
--
-- >>> let kdm = 'build' point3dAsList pointValuePairs
--
-- >>> 'nearest' kdm (Point3d 0.1 0.1 0.1)
-- [Point3d {x = 0.0, y = 0.0, z = 0.0}, \"First\"]
-- @

-- | A node of a /k/-d tree structure that stores a point of type @p@
-- with axis values of type @a@. Additionally, each point is
-- associated with a value of type @v@. Note: users typically will not
-- need to use this type, but we export it just in case.
data TreeNode a p v = TreeNode { forall a p v. TreeNode a p v -> TreeNode a p v
_treeLeft :: TreeNode a p v
                               , forall a p v. TreeNode a p v -> (p, v)
_treePoint :: (p, v)
                               , forall a p v. TreeNode a p v -> a
_axisValue :: a
                               , forall a p v. TreeNode a p v -> TreeNode a p v
_treeRight :: TreeNode a p v
                               } |
                      Empty
  deriving ((forall x. TreeNode a p v -> Rep (TreeNode a p v) x)
-> (forall x. Rep (TreeNode a p v) x -> TreeNode a p v)
-> Generic (TreeNode a p v)
forall x. Rep (TreeNode a p v) x -> TreeNode a p v
forall x. TreeNode a p v -> Rep (TreeNode 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 (TreeNode a p v) x -> TreeNode a p v
forall a p v x. TreeNode a p v -> Rep (TreeNode a p v) x
$cfrom :: forall a p v x. TreeNode a p v -> Rep (TreeNode a p v) x
from :: forall x. TreeNode a p v -> Rep (TreeNode a p v) x
$cto :: forall a p v x. Rep (TreeNode a p v) x -> TreeNode a p v
to :: forall x. Rep (TreeNode a p v) x -> TreeNode a p v
Generic, Int -> TreeNode a p v -> ShowS
[TreeNode a p v] -> ShowS
TreeNode a p v -> String
(Int -> TreeNode a p v -> ShowS)
-> (TreeNode a p v -> String)
-> ([TreeNode a p v] -> ShowS)
-> Show (TreeNode a p v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a p v.
(Show p, Show v, Show a) =>
Int -> TreeNode a p v -> ShowS
forall a p v. (Show p, Show v, Show a) => [TreeNode a p v] -> ShowS
forall a p v. (Show p, Show v, Show a) => TreeNode a p v -> String
$cshowsPrec :: forall a p v.
(Show p, Show v, Show a) =>
Int -> TreeNode a p v -> ShowS
showsPrec :: Int -> TreeNode a p v -> ShowS
$cshow :: forall a p v. (Show p, Show v, Show a) => TreeNode a p v -> String
show :: TreeNode a p v -> String
$cshowList :: forall a p v. (Show p, Show v, Show a) => [TreeNode a p v] -> ShowS
showList :: [TreeNode a p v] -> ShowS
Show, ReadPrec [TreeNode a p v]
ReadPrec (TreeNode a p v)
Int -> ReadS (TreeNode a p v)
ReadS [TreeNode a p v]
(Int -> ReadS (TreeNode a p v))
-> ReadS [TreeNode a p v]
-> ReadPrec (TreeNode a p v)
-> ReadPrec [TreeNode a p v]
-> Read (TreeNode a p v)
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall a p v. (Read p, Read v, Read a) => ReadPrec [TreeNode a p v]
forall a p v. (Read p, Read v, Read a) => ReadPrec (TreeNode a p v)
forall a p v.
(Read p, Read v, Read a) =>
Int -> ReadS (TreeNode a p v)
forall a p v. (Read p, Read v, Read a) => ReadS [TreeNode a p v]
$creadsPrec :: forall a p v.
(Read p, Read v, Read a) =>
Int -> ReadS (TreeNode a p v)
readsPrec :: Int -> ReadS (TreeNode a p v)
$creadList :: forall a p v. (Read p, Read v, Read a) => ReadS [TreeNode a p v]
readList :: ReadS [TreeNode a p v]
$creadPrec :: forall a p v. (Read p, Read v, Read a) => ReadPrec (TreeNode a p v)
readPrec :: ReadPrec (TreeNode a p v)
$creadListPrec :: forall a p v. (Read p, Read v, Read a) => ReadPrec [TreeNode a p v]
readListPrec :: ReadPrec [TreeNode a p v]
Read)
instance (NFData a, NFData p, NFData v) => NFData (TreeNode a p v) where rnf :: TreeNode a p v -> ()
rnf = TreeNode a p v -> ()
forall a. (Generic a, GNFData (Rep a)) => a -> ()
genericRnf

mapTreeNode :: (v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
mapTreeNode :: forall v1 v2 a p. (v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
mapTreeNode v1 -> v2
_ TreeNode a p v1
Empty = TreeNode a p v2
forall a p v. TreeNode a p v
Empty
mapTreeNode v1 -> v2
f (TreeNode TreeNode a p v1
left (p
k, v1
v) a
axisValue TreeNode a p v1
right) =
  TreeNode a p v2
-> (p, v2) -> a -> TreeNode a p v2 -> TreeNode a p v2
forall a p v.
TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
TreeNode ((v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
forall v1 v2 a p. (v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
mapTreeNode v1 -> v2
f TreeNode a p v1
left) (p
k, v1 -> v2
f v1
v) a
axisValue ((v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
forall v1 v2 a p. (v1 -> v2) -> TreeNode a p v1 -> TreeNode a p v2
mapTreeNode v1 -> v2
f TreeNode a p v1
right)

-- | Converts a point of type @p@ with axis values of type
-- @a@ into a list of axis values [a].
type PointAsListFn a p = p -> [a]

-- | Returns the squared distance between two points of type
-- @p@ with axis values of type @a@.
type SquaredDistanceFn a p = p -> p -> a

-- | A /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 { forall a p v. KdMap a p v -> PointAsListFn a p
_pointAsList :: PointAsListFn a p
                         , forall a p v. KdMap a p v -> SquaredDistanceFn a p
_distSqr     :: SquaredDistanceFn a p
                         , forall a p v. KdMap a p v -> TreeNode a p v
_rootNode    :: TreeNode a p v
                         , forall a p v. KdMap a p v -> Int
_size        :: 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
$cfrom :: forall a p v x. KdMap a p v -> Rep (KdMap a p v) x
from :: forall 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
to :: forall x. Rep (KdMap a p v) x -> KdMap a p v
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 PointAsListFn a p
_ SquaredDistanceFn a p
_ TreeNode a p v
rootNode Int
_) = String
"KdMap " String -> ShowS
forall a. [a] -> [a] -> [a]
++ TreeNode a p v -> String
forall a. Show a => a -> String
show TreeNode a p v
rootNode

instance Functor (KdMap a p) where
  fmap :: forall a b. (a -> b) -> KdMap a p a -> KdMap a p b
fmap a -> b
f KdMap a p a
kdMap = KdMap a p a
kdMap { _rootNode = mapTreeNode f (_rootNode kdMap) }

foldrTreeNode :: ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
foldrTreeNode :: forall p v b a. ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
foldrTreeNode (p, v) -> b -> b
_ b
z TreeNode a p v
Empty = b
z
foldrTreeNode (p, v) -> b -> b
f b
z (TreeNode TreeNode a p v
left (p, v)
p a
_ TreeNode a p v
right) =
  ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
forall p v b a. ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
foldrTreeNode (p, v) -> b -> b
f ((p, v) -> b -> b
f (p, v)
p (((p, v) -> b -> b) -> b -> TreeNode a p v -> b
forall p v b a. ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
foldrTreeNode (p, v) -> b -> b
f b
z TreeNode a p v
right)) TreeNode a p v
left

-- | Performs a foldr over each point-value pair in the 'KdMap'.
foldrWithKey :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey :: forall p v b a. ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey (p, v) -> b -> b
f b
z (KdMap PointAsListFn a p
_ SquaredDistanceFn a p
_ TreeNode a p v
r Int
_) = ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
forall p v b a. ((p, v) -> b -> b) -> b -> TreeNode a p v -> b
foldrTreeNode (p, v) -> b -> b
f b
z TreeNode a p v
r

instance Foldable (KdMap a p) where
  foldr :: forall a b. (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)

traverseTreeNode :: Applicative f => (b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
traverseTreeNode :: forall (f :: * -> *) b c a p.
Applicative f =>
(b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
traverseTreeNode b -> f c
_ TreeNode a p b
Empty = TreeNode a p c -> f (TreeNode a p c)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TreeNode a p c
forall a p v. TreeNode a p v
Empty
traverseTreeNode b -> f c
f (TreeNode TreeNode a p b
l (p
p, b
v) a
axisValue TreeNode a p b
r) =
  TreeNode a p c -> (p, c) -> a -> TreeNode a p c -> TreeNode a p c
forall a p v.
TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
TreeNode (TreeNode a p c -> (p, c) -> a -> TreeNode a p c -> TreeNode a p c)
-> f (TreeNode a p c)
-> f ((p, c) -> a -> TreeNode a p c -> TreeNode a p c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    (b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
forall (f :: * -> *) b c a p.
Applicative f =>
(b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
traverseTreeNode b -> f c
f TreeNode a p b
l f ((p, c) -> a -> TreeNode a p c -> TreeNode a p c)
-> f (p, c) -> f (a -> TreeNode a p c -> TreeNode a p c)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
    ((,) p
p (c -> (p, c)) -> f c -> f (p, c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> f c
f b
v) f (a -> TreeNode a p c -> TreeNode a p c)
-> f a -> f (TreeNode a p c -> TreeNode a p c)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> -- would simply be traverse f (p, v), but
                        -- base-4.6.* doesn't have a Traversable
                        -- instance for tuples.
    a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
axisValue f (TreeNode a p c -> TreeNode a p c)
-> f (TreeNode a p c) -> f (TreeNode a p c)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
    (b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
forall (f :: * -> *) b c a p.
Applicative f =>
(b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
traverseTreeNode b -> f c
f TreeNode a p b
r

instance Traversable (KdMap a p) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> KdMap a p a -> f (KdMap a p b)
traverse a -> f b
f (KdMap PointAsListFn a p
p SquaredDistanceFn a p
d TreeNode a p a
r Int
n) =
    PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p b -> Int -> KdMap a p b
forall a p v.
PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p v -> Int -> KdMap a p v
KdMap (PointAsListFn a p
 -> SquaredDistanceFn a p -> TreeNode a p b -> Int -> KdMap a p b)
-> f (PointAsListFn a p)
-> f (SquaredDistanceFn a p
      -> TreeNode a p b -> Int -> KdMap a p b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PointAsListFn a p -> f (PointAsListFn a p)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PointAsListFn a p
p f (SquaredDistanceFn a p -> TreeNode a p b -> Int -> KdMap a p b)
-> f (SquaredDistanceFn a p)
-> f (TreeNode a p b -> Int -> KdMap a p b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SquaredDistanceFn a p -> f (SquaredDistanceFn a p)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SquaredDistanceFn a p
d f (TreeNode a p b -> Int -> KdMap a p b)
-> f (TreeNode a p b) -> f (Int -> KdMap a p b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> TreeNode a p a -> f (TreeNode a p b)
forall (f :: * -> *) b c a p.
Applicative f =>
(b -> f c) -> TreeNode a p b -> f (TreeNode a p c)
traverseTreeNode a -> f b
f TreeNode a p a
r f (Int -> KdMap a p b) -> f Int -> f (KdMap a p b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> f Int
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n

-- | Builds an empty 'KdMap'.
empty :: Real a => PointAsListFn a p -> KdMap a p v
empty :: forall a p v. Real a => 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.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist PointAsListFn a p
p2l (PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l)

-- | Builds an empty 'KdMap' using a user-specified squared distance
-- function.
emptyWithDist :: Real a => PointAsListFn a p
                        -> SquaredDistanceFn a p
                        -> KdMap a p v
emptyWithDist :: forall a p v.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 = PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p v -> Int -> KdMap a p v
forall a p v.
PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p v -> Int -> KdMap a p v
KdMap PointAsListFn a p
p2l SquaredDistanceFn a p
d2 TreeNode a p v
forall a p v. TreeNode a p v
Empty Int
0

-- | Returns 'True' if the given 'KdMap' is empty.
null :: KdMap a p v -> Bool
null :: forall a p a. KdMap a p a -> Bool
null KdMap a p v
kdm = KdMap a p v -> Int
forall a p v. KdMap a p v -> Int
_size KdMap a p v
kdm Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0

-- | Builds a 'KdMap' with a single point-value pair and a
-- user-specified squared distance function.
singletonWithDist :: Real a => PointAsListFn a p
                            -> SquaredDistanceFn a p
                            -> (p, v)
                            -> KdMap a p v
singletonWithDist :: 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
d2 (p
p, v
v) =
  let singletonTreeNode :: TreeNode a p v
singletonTreeNode = TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
forall a p v.
TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
TreeNode TreeNode a p v
forall a p v. TreeNode a p v
Empty (p
p, v
v) ([a] -> a
forall a. HasCallStack => [a] -> a
head ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p
p2l p
p) TreeNode a p v
forall a p v. TreeNode a p v
Empty
  in  PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p v -> Int -> KdMap a p v
forall a p v.
PointAsListFn a p
-> SquaredDistanceFn a p -> TreeNode a p v -> Int -> KdMap a p v
KdMap PointAsListFn a p
p2l SquaredDistanceFn a p
d2 TreeNode a p v
singletonTreeNode Int
1

-- | Builds a 'KdMap' with a single point-value pair.
singleton :: Real a => PointAsListFn a p -> (p, v) -> KdMap a p v
singleton :: forall a p v. Real a => PointAsListFn a p -> (p, v) -> KdMap a p v
singleton PointAsListFn a p
p2l (p
p, v
v) = 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 (PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l) (p
p, v
v)

quickselect :: (b -> b -> Ordering) -> Int -> [b] -> b
quickselect :: forall b. (b -> b -> Ordering) -> Int -> [b] -> b
quickselect b -> b -> Ordering
cmp = Int -> [b] -> b
go
  where go :: Int -> [b] -> b
go Int
_ [] = String -> b
forall a. HasCallStack => String -> a
error String
"quickselect must be called on a non-empty list."
        go Int
k (b
x:[b]
xs) | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = Int -> [b] -> b
go Int
k [b]
ys
                    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
l = Int -> [b] -> b
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [b]
zs
                    | Bool
otherwise = b
x
          where ([b]
ys, [b]
zs) = (b -> Bool) -> [b] -> ([b], [b])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition ((Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
LT) (Ordering -> Bool) -> (b -> Ordering) -> b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> b -> Ordering
`cmp` b
x)) [b]
xs
                l :: Int
l = [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
ys

-- | Builds a 'KdMap' from a list of pairs of points (of type p) and
-- values (of type v), using a user-specified squared distance
-- function.
--
-- Average time complexity: /O(n * log(n))/ for /n/ data points.
--
-- Worst case time complexity: /O(n^2)/ for /n/ data points.
--
-- Worst case space complexity: /O(n)/ for /n/ data points.
buildWithDist :: Real a => PointAsListFn a p
                        -> SquaredDistanceFn a p
                        -> [(p, v)]
                        -> KdMap a p v
buildWithDist :: forall a p v.
Real a =>
PointAsListFn a p
-> SquaredDistanceFn a p -> [(p, v)] -> KdMap a p v
buildWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 [] = PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
forall a p v.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2
buildWithDist PointAsListFn a p
pointAsList SquaredDistanceFn a p
distSqr [(p, v)]
dataPoints =
  let axisValsPointsPairs :: [([a], (p, v))]
axisValsPointsPairs = [[a]] -> [(p, v)] -> [([a], (p, v))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((p, v) -> [a]) -> [(p, v)] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle ([a] -> [a]) -> ((p, v) -> [a]) -> (p, v) -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PointAsListFn a p
pointAsList PointAsListFn a p -> ((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)]
dataPoints) [(p, v)]
dataPoints
  in  KdMap { _pointAsList :: PointAsListFn a p
_pointAsList = PointAsListFn a p
pointAsList
            , _distSqr :: SquaredDistanceFn a p
_distSqr     = SquaredDistanceFn a p
distSqr
            , _rootNode :: TreeNode a p v
_rootNode    = [([a], (p, v))] -> TreeNode a p v
forall {a} {p} {v}. Ord a => [([a], (p, v))] -> TreeNode a p v
buildTreeInternal [([a], (p, v))]
axisValsPointsPairs
            , _size :: Int
_size        = [(p, v)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(p, v)]
dataPoints
            }
  where buildTreeInternal :: [([a], (p, v))] -> TreeNode a p v
buildTreeInternal [] = TreeNode a p v
forall a p v. TreeNode a p v
Empty
        buildTreeInternal [([a], (p, v))]
ps =
          let n :: Int
n = [([a], (p, v))] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [([a], (p, v))]
ps
              (a
medianAxisVal : [a]
_, (p, v)
_) =
                (([a], (p, v)) -> ([a], (p, v)) -> Ordering)
-> Int -> [([a], (p, v))] -> ([a], (p, v))
forall b. (b -> b -> Ordering) -> Int -> [b] -> b
quickselect ((([a], (p, v)) -> a) -> ([a], (p, v)) -> ([a], (p, v)) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing ([a] -> a
forall a. HasCallStack => [a] -> a
head ([a] -> a) -> (([a], (p, v)) -> [a]) -> ([a], (p, v)) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a], (p, v)) -> [a]
forall a b. (a, b) -> a
fst)) (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [([a], (p, v))]
ps
              f :: ([a], a)
-> ([([a], a)], Maybe a, [([a], a)])
-> ([([a], a)], Maybe a, [([a], a)])
f ([], a
_) ([([a], a)], Maybe a, [([a], a)])
_ = String -> ([([a], a)], Maybe a, [([a], a)])
forall a. HasCallStack => String -> a
error String
"buildKdMap.f: no empty lists allowed!"
              f (a
v : [a]
vt, a
p) ([([a], a)]
lt, Maybe a
maybeMedian, [([a], a)]
gt)
                | a
v a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
medianAxisVal = (([a]
vt, a
p) ([a], a) -> [([a], a)] -> [([a], a)]
forall a. a -> [a] -> [a]
: [([a], a)]
lt, Maybe a
maybeMedian, [([a], a)]
gt)
                | a
v a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
medianAxisVal = ([([a], a)]
lt, Maybe a
maybeMedian, ([a]
vt, a
p) ([a], a) -> [([a], a)] -> [([a], a)]
forall a. a -> [a] -> [a]
: [([a], a)]
gt)
                | Bool
otherwise =
                    case Maybe a
maybeMedian of
                      Maybe a
Nothing -> ([([a], a)]
lt, a -> Maybe a
forall a. a -> Maybe a
Just a
p, [([a], a)]
gt)
                      Just a
_ -> (([a]
vt, a
p) ([a], a) -> [([a], a)] -> [([a], a)]
forall a. a -> [a] -> [a]
: [([a], a)]
lt, Maybe a
maybeMedian, [([a], a)]
gt)
              ([([a], (p, v))]
leftPoints, Maybe (p, v)
maybeMedianPt, [([a], (p, v))]
rightPoints) = (([a], (p, v))
 -> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))])
 -> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))]))
-> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))])
-> [([a], (p, v))]
-> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))])
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
L.foldr ([a], (p, v))
-> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))])
-> ([([a], (p, v))], Maybe (p, v), [([a], (p, v))])
forall {a}.
([a], a)
-> ([([a], a)], Maybe a, [([a], a)])
-> ([([a], a)], Maybe a, [([a], a)])
f ([], Maybe (p, v)
forall a. Maybe a
Nothing, []) [([a], (p, v))]
ps
          in  TreeNode
              { _treeLeft :: TreeNode a p v
_treeLeft  = [([a], (p, v))] -> TreeNode a p v
buildTreeInternal [([a], (p, v))]
leftPoints
              , _treePoint :: (p, v)
_treePoint = Maybe (p, v) -> (p, v)
forall a. HasCallStack => Maybe a -> a
fromJust Maybe (p, v)
maybeMedianPt
              , _axisValue :: a
_axisValue = a
medianAxisVal
              , _treeRight :: TreeNode a p v
_treeRight = [([a], (p, v))] -> TreeNode a p v
buildTreeInternal [([a], (p, v))]
rightPoints
              }

-- | A default implementation of squared distance given two points and
-- a 'PointAsListFn'.
defaultSqrDist :: Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist :: forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
pointAsList p
k1 p
k2 =
  [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
L.sum ([a] -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) (PointAsListFn a p
pointAsList p
k1) (PointAsListFn a p
pointAsList p
k2)

-- | Builds a 'KdTree' from a list of pairs of points (of type p) and
-- values (of type v) using a default squared distance function
-- 'defaultSqrDist'.
--
-- Average complexity: /O(n * log(n))/ for /n/ data points.
--
-- Worst case time complexity: /O(n^2)/ for /n/ data points.
--
-- Worst case space complexity: /O(n)/ for /n/ data points.
build :: Real a => PointAsListFn a p -> [(p, v)] -> KdMap a p v
build :: forall a p v.
Real a =>
PointAsListFn a p -> [(p, v)] -> KdMap a p v
build PointAsListFn a p
pointAsList =
  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
buildWithDist PointAsListFn a p
pointAsList (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
pointAsList

-- | Inserts a point-value pair into a 'KdMap'. This can potentially
-- cause the internal tree structure to become unbalanced. If the tree
-- becomes too unbalanced, point queries will be very inefficient. If
-- you need to perform lots of point insertions on an already existing
-- /k/-d map, check out
-- @Data.KdMap.Dynamic.@'Data.KdMap.Dynamic.KdMap'.
--
-- Average complexity: /O(log(n))/ for /n/ data points.
--
-- Worst case time complexity: /O(n)/ for /n/ data points.
insertUnbalanced :: Real a => KdMap a p v -> p -> v -> KdMap a p v
insertUnbalanced :: forall a p v. Real a => KdMap a p v -> p -> v -> KdMap a p v
insertUnbalanced kdm :: KdMap a p v
kdm@(KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
_ TreeNode a p v
rootNode Int
n) p
p' v
v' =
  KdMap a p v
kdm { _rootNode = go rootNode (cycle $ pointAsList p'), _size = n + 1 }
  where
    go :: TreeNode a p v -> [a] -> TreeNode a p v
go TreeNode a p v
_ [] = String -> TreeNode a p v
forall a. HasCallStack => String -> a
error String
"insertUnbalanced.go: no empty lists allowed!"
    go TreeNode a p v
Empty (a
axisValue' : [a]
_) = TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
forall a p v.
TreeNode a p v -> (p, v) -> a -> TreeNode a p v -> TreeNode a p v
TreeNode TreeNode a p v
forall a p v. TreeNode a p v
Empty (p
p', v
v') a
axisValue' TreeNode a p v
forall a p v. TreeNode a p v
Empty
    go t :: TreeNode a p v
t@(TreeNode TreeNode a p v
left (p, v)
_ a
nodeAxisValue TreeNode a p v
right) (a
axisValue' : [a]
nextAxisValues)
      | a
axisValue' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nodeAxisValue = TreeNode a p v
t { _treeLeft = go left nextAxisValues }
      | Bool
otherwise = TreeNode a p v
t { _treeRight = go right nextAxisValues }

-- | Inserts a list of point-value pairs into a 'KdMap'. This can
-- potentially cause the internal tree structure to become unbalanced,
-- which leads to inefficient point queries.
--
-- Average complexity: /O(n * log(n))/ for /n/ data points.
--
-- Worst case time complexity: /O(n^2)/ for /n/ data points.
batchInsertUnbalanced :: Real a => KdMap a p v -> [(p, v)] -> KdMap a p v
batchInsertUnbalanced :: forall a p v. Real a => KdMap a p v -> [(p, v)] -> KdMap a p v
batchInsertUnbalanced = (KdMap a p v -> (p, v) -> KdMap a p v)
-> KdMap a p v -> [(p, v)] -> KdMap a p v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((KdMap a p v -> (p, v) -> KdMap a p v)
 -> KdMap a p v -> [(p, v)] -> KdMap a p v)
-> (KdMap a p v -> (p, v) -> KdMap a p v)
-> KdMap a p v
-> [(p, v)]
-> KdMap a p v
forall a b. (a -> b) -> a -> b
$ \KdMap a p v
kdm (p
p, v
v) -> 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
insertUnbalanced KdMap a p v
kdm p
p v
v

assocsInternal :: TreeNode a p v -> [(p, v)]
assocsInternal :: forall a p v. TreeNode a p v -> [(p, v)]
assocsInternal TreeNode a p v
t = TreeNode a p v -> [(p, v)] -> [(p, v)]
forall {a} {p} {v}. TreeNode a p v -> [(p, v)] -> [(p, v)]
go TreeNode a p v
t []
  where go :: TreeNode a p v -> [(p, v)] -> [(p, v)]
go TreeNode a p v
Empty = [(p, v)] -> [(p, v)]
forall a. a -> a
id
        go (TreeNode TreeNode a p v
l (p, v)
p a
_ TreeNode a p v
r) = TreeNode a p v -> [(p, v)] -> [(p, v)]
go TreeNode a p v
l ([(p, v)] -> [(p, v)])
-> ([(p, v)] -> [(p, v)]) -> [(p, v)] -> [(p, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((p, v)
p (p, v) -> [(p, v)] -> [(p, v)]
forall a. a -> [a] -> [a]
:) ([(p, v)] -> [(p, v)])
-> ([(p, v)] -> [(p, v)]) -> [(p, v)] -> [(p, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeNode a p v -> [(p, v)] -> [(p, v)]
go TreeNode a p v
r

-- | 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 :: forall a p v. KdMap a p v -> [(p, v)]
assocs (KdMap PointAsListFn a p
_ SquaredDistanceFn a p
_ TreeNode a p v
t Int
_) = TreeNode a p v -> [(p, v)]
forall a p v. TreeNode a p v -> [(p, v)]
assocsInternal TreeNode a p v
t

-- | Returns all points in the 'KdMap'.
--
-- Time complexity: /O(n)/ for /n/ data points.
keys :: KdMap a p v -> [p]
keys :: forall a p v. 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 :: forall a p a. KdMap a p a -> [a]
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

-- | 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(n))/ for /n/ data points.
--
-- Worst case time complexity: /O(n)/ for /n/ data points.
--
-- Throws error if called on an empty 'KdMap'.
nearest :: Real a => KdMap a p v -> p -> (p, v)
nearest :: forall a p v. Real a => KdMap a p v -> p -> (p, v)
nearest (KdMap PointAsListFn a p
_ SquaredDistanceFn a p
_ TreeNode a p v
Empty Int
_) p
_ =
  String -> (p, v)
forall a. HasCallStack => String -> a
error String
"Attempted to call nearest on an empty KdMap."
nearest (KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
distSqr t :: TreeNode a p v
t@(TreeNode TreeNode a p v
_ (p, v)
root a
_ TreeNode a p v
_) Int
_) p
query =
  -- This is an ugly way to kickstart the function but it's faster
  -- than using a Maybe.
  ((p, v), a) -> (p, v)
forall a b. (a, b) -> a
fst (((p, v), a) -> (p, v)) -> ((p, v), a) -> (p, v)
forall a b. (a -> b) -> a -> b
$ ((p, v), a) -> [a] -> TreeNode a p v -> ((p, v), a)
forall {v}. ((p, v), a) -> [a] -> TreeNode a p v -> ((p, v), a)
go ((p, v)
root, SquaredDistanceFn a p
distSqr ((p, v) -> p
forall a b. (a, b) -> a
fst (p, v)
root) p
query) ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p
pointAsList p
query) TreeNode a p v
t
  where
    go :: ((p, v), a) -> [a] -> TreeNode a p v -> ((p, v), a)
go ((p, v), a)
_ [] TreeNode a p v
_ = String -> ((p, v), a)
forall a. HasCallStack => String -> a
error String
"nearest.go: no empty lists allowed!"
    go ((p, v), a)
bestSoFar [a]
_ TreeNode a p v
Empty = ((p, v), a)
bestSoFar
    go ((p, v), a)
bestSoFar
       (a
queryAxisValue : [a]
qvs)
       (TreeNode TreeNode a p v
left (p
nodeK, v
nodeV) a
nodeAxisVal TreeNode a p v
right) =
      let better :: (a, a) -> (a, a) -> (a, a)
better x1 :: (a, a)
x1@(a
_, a
dist1) x2 :: (a, a)
x2@(a
_, a
dist2) = if a
dist1 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
dist2
                                               then (a, a)
x1
                                               else (a, a)
x2
          currDist :: a
currDist       = SquaredDistanceFn a p
distSqr p
query p
nodeK
          bestAfterNode :: ((p, v), a)
bestAfterNode = ((p, v), a) -> ((p, v), a) -> ((p, v), a)
forall {a} {a}. Ord a => (a, a) -> (a, a) -> (a, a)
better ((p
nodeK, v
nodeV), a
currDist) ((p, v), a)
bestSoFar
          nearestInTree :: TreeNode a p v -> TreeNode a p v -> ((p, v), a)
nearestInTree TreeNode a p v
onsideSubtree TreeNode a p v
offsideSubtree =
            let bestAfterOnside :: ((p, v), a)
bestAfterOnside = ((p, v), a) -> [a] -> TreeNode a p v -> ((p, v), a)
go ((p, v), a)
bestAfterNode [a]
qvs TreeNode a p v
onsideSubtree
                checkOffsideSubtree :: Bool
checkOffsideSubtree =
                  (a
queryAxisValue a -> a -> a
forall a. Num a => a -> a -> a
- a
nodeAxisVal)a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2 :: Int) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< ((p, v), a) -> a
forall a b. (a, b) -> b
snd ((p, v), a)
bestAfterOnside
            in  if Bool
checkOffsideSubtree
                then ((p, v), a) -> [a] -> TreeNode a p v -> ((p, v), a)
go ((p, v), a)
bestAfterOnside [a]
qvs TreeNode a p v
offsideSubtree
                else ((p, v), a)
bestAfterOnside
      in  if a
queryAxisValue a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nodeAxisVal
          then TreeNode a p v -> TreeNode a p v -> ((p, v), a)
nearestInTree TreeNode a p v
left TreeNode a p v
right
          else TreeNode a p v -> TreeNode a p v -> ((p, v), a)
nearestInTree TreeNode a p v
right TreeNode a p v
left

-- | Given a 'KdMap', a query point, and a radius, returns all
-- point-value pairs in the 'KdMap' 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 and a radius
-- that spans all points in the structure.
inRadius :: Real a => KdMap a p v
                   -> a -- ^ radius
                   -> p -- ^ query point
                   -> [(p, v)] -- ^ list of point-value pairs with
                               -- points within given radius of query
inRadius :: forall a p v. Real a => KdMap a p v -> a -> p -> [(p, v)]
inRadius (KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
distSqr TreeNode a p v
t Int
_) a
radius p
query =
  [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
forall {v}. [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p
pointAsList p
query) TreeNode a p v
t []
  where
    go :: [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go [] TreeNode a p v
_ [(p, v)]
_ = String -> [(p, v)]
forall a. HasCallStack => String -> a
error String
"inRadius.go: no empty lists allowed!"
    go [a]
_ TreeNode a p v
Empty [(p, v)]
acc = [(p, v)]
acc
    go (a
queryAxisValue : [a]
qvs) (TreeNode TreeNode a p v
left (p
k, v
v) a
nodeAxisVal TreeNode a p v
right) [(p, v)]
acc =
      let onTheLeft :: Bool
onTheLeft = a
queryAxisValue a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nodeAxisVal
          accAfterOnside :: [(p, v)]
accAfterOnside = if   Bool
onTheLeft
                           then [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go [a]
qvs TreeNode a p v
left [(p, v)]
acc
                           else [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go [a]
qvs TreeNode a p v
right [(p, v)]
acc
          accAfterOffside :: [(p, v)]
accAfterOffside = if   a -> a
forall a. Num a => a -> a
abs (a
queryAxisValue a -> a -> a
forall a. Num a => a -> a -> a
- a
nodeAxisVal) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
radius
                            then if   Bool
onTheLeft
                                 then [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go [a]
qvs TreeNode a p v
right [(p, v)]
accAfterOnside
                                 else [a] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
go [a]
qvs TreeNode a p v
left [(p, v)]
accAfterOnside
                            else [(p, v)]
accAfterOnside
          accAfterCurrent :: [(p, v)]
accAfterCurrent = if SquaredDistanceFn a p
distSqr p
k p
query a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
radius a -> a -> a
forall a. Num a => a -> a -> a
* a
radius
                            then (p
k, v
v) (p, v) -> [(p, v)] -> [(p, v)]
forall a. a -> [a] -> [a]
: [(p, v)]
accAfterOffside
                            else [(p, v)]
accAfterOffside
      in  [(p, v)]
accAfterCurrent

-- | 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(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 :: forall a p v. Real a => KdMap a p v -> Int -> p -> [(p, v)]
kNearest (KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
distSqr TreeNode a p v
t Int
_) Int
numNeighbors p
query =
  [(p, v)] -> [(p, v)]
forall a. [a] -> [a]
reverse ([(p, v)] -> [(p, v)]) -> [(p, v)] -> [(p, v)]
forall a b. (a -> b) -> a -> b
$ ((a, (p, v)) -> (p, v)) -> [(a, (p, v))] -> [(p, v)]
forall a b. (a -> b) -> [a] -> [b]
map (a, (p, v)) -> (p, v)
forall a b. (a, b) -> b
snd ([(a, (p, v))] -> [(p, v)]) -> [(a, (p, v))] -> [(p, v)]
forall a b. (a -> b) -> a -> b
$ Heap FstMaxPolicy (a, (p, v)) -> [(a, (p, v))]
forall pol item. HeapItem pol item => Heap pol item -> [item]
Q.toAscList (Heap FstMaxPolicy (a, (p, v)) -> [(a, (p, v))])
-> Heap FstMaxPolicy (a, (p, v)) -> [(a, (p, v))]
forall a b. (a -> b) -> a -> b
$ [a]
-> Heap FstMaxPolicy (a, (p, v))
-> TreeNode a p v
-> Heap FstMaxPolicy (a, (p, v))
forall {pol} {v}.
HeapItem pol (a, (p, v)) =>
[a]
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
go ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p
pointAsList p
query)
    (HeapT (Prio FstMaxPolicy (a, (p, v))) (p, v)
HeapT
  (Prio FstMaxPolicy (a, (p, v))) (Val FstMaxPolicy (a, (p, v)))
forall prio val. HeapT prio val
forall {a} {p} {v}.
HeapT
  (Prio FstMaxPolicy (a, (p, v))) (Val FstMaxPolicy (a, (p, v)))
Q.empty :: Q.MaxPrioHeap a (p,v)) TreeNode a p v
t
  where
    -- go :: [a] -> Q.MaxPrioHeap a (p, v) -> TreeNode a p v -> Q.MaxPrioHeap a (p, v)
    go :: [a]
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
go [] HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
_ TreeNode a p v
_ = String -> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
forall a. HasCallStack => String -> a
error String
"kNearest.go: no empty lists allowed!"
    go [a]
_ HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q TreeNode a p v
Empty = HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q
    go (a
queryAxisValue : [a]
qvs) HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q (TreeNode TreeNode a p v
left (p
k, v
v) a
nodeAxisVal TreeNode a p v
right) =
      let insertBounded :: HeapT (Prio pol (a, b)) (Val pol (a, b))
-> a -> b -> HeapT (Prio pol (a, b)) (Val pol (a, b))
insertBounded HeapT (Prio pol (a, b)) (Val pol (a, b))
queue a
dist b
x
            | HeapT (Prio pol (a, b)) (Val pol (a, b)) -> Int
forall prio val. HeapT prio val -> Int
Q.size HeapT (Prio pol (a, b)) (Val pol (a, b))
queue Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
numNeighbors = (a, b)
-> HeapT (Prio pol (a, b)) (Val pol (a, b))
-> HeapT (Prio pol (a, b)) (Val pol (a, b))
forall pol item.
HeapItem pol item =>
item -> Heap pol item -> Heap pol item
Q.insert (a
dist, b
x) HeapT (Prio pol (a, b)) (Val pol (a, b))
queue
            | Bool
otherwise = let ((a
farthestDist, b
_), HeapT (Prio pol (a, b)) (Val pol (a, b))
rest) = Maybe ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
-> ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
 -> ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b))))
-> Maybe ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
-> ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
forall a b. (a -> b) -> a -> b
$ HeapT (Prio pol (a, b)) (Val pol (a, b))
-> Maybe ((a, b), HeapT (Prio pol (a, b)) (Val pol (a, b)))
forall pol item.
HeapItem pol item =>
Heap pol item -> Maybe (item, Heap pol item)
Q.view HeapT (Prio pol (a, b)) (Val pol (a, b))
queue
                          in  if a
dist a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
farthestDist
                              then (a, b)
-> HeapT (Prio pol (a, b)) (Val pol (a, b))
-> HeapT (Prio pol (a, b)) (Val pol (a, b))
forall pol item.
HeapItem pol item =>
item -> Heap pol item -> Heap pol item
Q.insert (a
dist, b
x) HeapT (Prio pol (a, b)) (Val pol (a, b))
rest
                              else HeapT (Prio pol (a, b)) (Val pol (a, b))
queue
          q' :: HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q' = HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> a
-> (p, v)
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
forall {pol} {a} {b}.
(HeapItem pol (a, b), Ord a) =>
HeapT (Prio pol (a, b)) (Val pol (a, b))
-> a -> b -> HeapT (Prio pol (a, b)) (Val pol (a, b))
insertBounded HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q (SquaredDistanceFn a p
distSqr p
k p
query) (p
k, v
v)
          kNear :: HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
kNear HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue TreeNode a p v
onsideSubtree TreeNode a p v
offsideSubtree =
            let queue' :: HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue' = [a]
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
go [a]
qvs HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue TreeNode a p v
onsideSubtree
                checkOffsideTree :: Bool
checkOffsideTree =
                  HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))) -> Int
forall prio val. HeapT prio val -> Int
Q.size HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
numNeighbors Bool -> Bool -> Bool
||
                  (a
queryAxisValue a -> a -> a
forall a. Num a => a -> a -> a
- a
nodeAxisVal)a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2 :: Int) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<
                    ((a, (p, v)) -> a
forall a b. (a, b) -> a
fst ((a, (p, v)) -> a)
-> (((a, (p, v)),
     HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
    -> (a, (p, v)))
-> ((a, (p, v)),
    HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, (p, v)), HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
-> (a, (p, v))
forall a b. (a, b) -> a
fst) (Maybe
  ((a, (p, v)), HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
-> ((a, (p, v)),
    HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe
   ((a, (p, v)), HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
 -> ((a, (p, v)),
     HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))))
-> Maybe
     ((a, (p, v)), HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
-> ((a, (p, v)),
    HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
forall a b. (a -> b) -> a -> b
$ HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> Maybe
     ((a, (p, v)), HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v))))
forall pol item.
HeapItem pol item =>
Heap pol item -> Maybe (item, Heap pol item)
Q.view HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue')
            in  if Bool
checkOffsideTree
                then [a]
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
go [a]
qvs HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue' TreeNode a p v
offsideSubtree
                else HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
queue'
      in  if a
queryAxisValue a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nodeAxisVal
          then HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
kNear HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q' TreeNode a p v
left TreeNode a p v
right
          else HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
-> TreeNode a p v
-> TreeNode a p v
-> HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
kNear HeapT (Prio pol (a, (p, v))) (Val pol (a, (p, v)))
q' TreeNode a p v
right TreeNode a p v
left

-- | 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.
--
-- TODO: Maybe use known bounds on entire tree structure to be able to
-- automatically count whole portions of tree as being within given
-- range.
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 :: forall a p v. Real a => KdMap a p v -> p -> p -> [(p, v)]
inRange (KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
_ TreeNode a p v
t Int
_) p
lowers p
uppers =
  [(a, a)] -> TreeNode a p v -> [(p, v)] -> [(p, v)]
forall {b} {v}.
Ord b =>
[(b, b)] -> TreeNode b p v -> [(p, v)] -> [(p, v)]
go ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle (PointAsListFn a p
pointAsList p
lowers) [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle (PointAsListFn a p
pointAsList p
uppers)) TreeNode a p v
t []
  where
    go :: [(b, b)] -> TreeNode b p v -> [(p, v)] -> [(p, v)]
go [] TreeNode b p v
_ [(p, v)]
_ = String -> [(p, v)]
forall a. HasCallStack => String -> a
error String
"inRange.go: no empty lists allowed!"
    go [(b, b)]
_ TreeNode b p v
Empty [(p, v)]
acc = [(p, v)]
acc
    go ((b
lower, b
upper) : [(b, b)]
nextBounds) (TreeNode TreeNode b p v
left (p, v)
p b
nodeAxisVal TreeNode b p v
right) [(p, v)]
acc =
      let accAfterLeft :: [(p, v)]
accAfterLeft = if b
lower b -> b -> Bool
forall a. Ord a => a -> a -> Bool
<= b
nodeAxisVal
                         then [(b, b)] -> TreeNode b p v -> [(p, v)] -> [(p, v)]
go [(b, b)]
nextBounds TreeNode b p v
left [(p, v)]
acc
                         else [(p, v)]
acc
          accAfterRight :: [(p, v)]
accAfterRight = if b
upper b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
nodeAxisVal
                          then [(b, b)] -> TreeNode b p v -> [(p, v)] -> [(p, v)]
go [(b, b)]
nextBounds TreeNode b p v
right [(p, v)]
accAfterLeft
                          else [(p, v)]
accAfterLeft
          valInRange :: a -> a -> a -> Bool
valInRange a
l a
x a
u = a
l a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
x Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
u
          -- maybe "cache" lowers and uppers as lists sooner as hint
          -- to ghc. Also, maybe only need to check previously
          -- unchecked axes?
          currentInRange :: Bool
currentInRange =
            [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
L.and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (a -> a -> a -> Bool) -> [a] -> [a] -> [a] -> [Bool]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 a -> a -> a -> Bool
forall {a}. Ord a => a -> a -> a -> Bool
valInRange
              (PointAsListFn a p
pointAsList p
lowers) (PointAsListFn a p
pointAsList PointAsListFn a p -> PointAsListFn a p
forall a b. (a -> b) -> a -> b
$ (p, v) -> p
forall a b. (a, b) -> a
fst (p, v)
p) (PointAsListFn a p
pointAsList p
uppers)
          accAfterCurrent :: [(p, v)]
accAfterCurrent = if Bool
currentInRange
                            then (p, v)
p (p, v) -> [(p, v)] -> [(p, v)]
forall a. a -> [a] -> [a]
: [(p, v)]
accAfterRight
                            else [(p, v)]
accAfterRight
      in  [(p, v)]
accAfterCurrent

-- | Returns the number of point-value pairs in the 'KdMap'.
--
-- Time complexity: /O(1)/
size :: KdMap a p v -> Int
size :: forall a p v. KdMap a p v -> Int
size (KdMap PointAsListFn a p
_ SquaredDistanceFn a p
_ TreeNode a p v
_ Int
n) = Int
n

isTreeNodeValid :: Real a => PointAsListFn a p -> Int -> TreeNode a p v -> Bool
isTreeNodeValid :: forall a p v.
Real a =>
PointAsListFn a p -> Int -> TreeNode a p v -> Bool
isTreeNodeValid PointAsListFn a p
_ Int
_ TreeNode a p v
Empty = Bool
True
isTreeNodeValid PointAsListFn a p
pointAsList Int
axis (TreeNode TreeNode a p v
l (p
k, v
_) a
nodeAxisVal TreeNode a p v
r) =
  let childrenAxisValues :: TreeNode a p b -> [a]
childrenAxisValues = ((p, b) -> a) -> [(p, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (([a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) ([a] -> a) -> ((p, b) -> [a]) -> (p, b) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PointAsListFn a p
pointAsList PointAsListFn a p -> ((p, b) -> p) -> (p, b) -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (p, b) -> p
forall a b. (a, b) -> a
fst) ([(p, b)] -> [a])
-> (TreeNode a p b -> [(p, b)]) -> TreeNode a p b -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeNode a p b -> [(p, b)]
forall a p v. TreeNode a p v -> [(p, v)]
assocsInternal
      leftSubtreeLess :: Bool
leftSubtreeLess = (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
L.all (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nodeAxisVal) ([a] -> Bool) -> [a] -> Bool
forall a b. (a -> b) -> a -> b
$ TreeNode a p v -> [a]
forall {a} {b}. TreeNode a p b -> [a]
childrenAxisValues TreeNode a p v
l
      rightSubtreeGreater :: Bool
rightSubtreeGreater = (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
L.all (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
nodeAxisVal) ([a] -> Bool) -> [a] -> Bool
forall a b. (a -> b) -> a -> b
$ TreeNode a p v -> [a]
forall {a} {b}. TreeNode a p b -> [a]
childrenAxisValues TreeNode a p v
r
      nextAxis :: Int
nextAxis = (Int
axis Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (PointAsListFn a p
pointAsList p
k)
  in  Bool
leftSubtreeLess Bool -> Bool -> Bool
&& Bool
rightSubtreeGreater Bool -> Bool -> Bool
&&
      PointAsListFn a p -> Int -> TreeNode a p v -> Bool
forall a p v.
Real a =>
PointAsListFn a p -> Int -> TreeNode a p v -> Bool
isTreeNodeValid PointAsListFn a p
pointAsList Int
nextAxis TreeNode a p v
l Bool -> Bool -> Bool
&& PointAsListFn a p -> Int -> TreeNode a p v -> Bool
forall a p v.
Real a =>
PointAsListFn a p -> Int -> TreeNode a p v -> Bool
isTreeNodeValid PointAsListFn a p
pointAsList Int
nextAxis TreeNode a p v
r

-- | Returns 'True' if tree structure adheres to k-d tree
-- properties. For internal testing use.
isValid :: Real a => KdMap a p v -> Bool
isValid :: forall a p v. Real a => KdMap a p v -> Bool
isValid (KdMap PointAsListFn a p
pointAsList SquaredDistanceFn a p
_ TreeNode a p v
r Int
_) = PointAsListFn a p -> Int -> TreeNode a p v -> Bool
forall a p v.
Real a =>
PointAsListFn a p -> Int -> TreeNode a p v -> Bool
isTreeNodeValid PointAsListFn a p
pointAsList Int
0 TreeNode a p v
r