module Data.PKTree (Point, Rectangle, PKTree, Node(..), cell, pointCell, rect, insert, radiusSearch) where
import Data.Tree
import Data.Maybe
import Data.List hiding (insert)
type Point = [Float]
type Rectangle = (Point, Point)
data Node a = Inner Rectangle | Leaf Point a deriving (Eq, Show, Read)
type PKTree a = Tree (Node a)
npartition :: (Integral b) => (a -> b) -> [a] -> [[a]]
npartition f =
map (map snd) . groupBy (onFst (==)) . sortBy (onFst compare) . map (\x -> (f x, x))
where
onFst f (a,_) (b,_) = f a b
rectContains :: Rectangle -> Rectangle -> Bool
rectContains (l, u) (l', u') =
all (uncurry (<=)) (zip l l') && all (uncurry (>)) (zip u u')
cell :: Rectangle -> PKTree a
cell rect = Node {
rootLabel = Inner rect,
subForest = []
}
pointCell :: Point -> a -> PKTree a
pointCell p v = Node {
rootLabel = Leaf p v,
subForest = []
}
rect :: PKTree a -> Rectangle
rect =
rect' . rootLabel
where
rect' (Inner r) = r
rect' (Leaf p _) = (p,p)
insert :: (Eq a) =>
Int
-> [Int]
-> PKTree a
-> Point
-> a
-> PKTree a
insert k r n@(Node {rootLabel = lbl, subForest = children}) p v =
Node {
rootLabel = lbl,
subForest = instantiateDivisions k r (rect n) newKids
}
where
newKids = insert' k r children p v
insert' :: (Eq a) => Int -> [Int] -> [PKTree a] -> Point -> a -> [PKTree a]
insert' k r children p v
| children `contains` p =
let (Just c) = children `maybecontain` p in
let c' = insert k r c p v in
(if length (subForest c') < k then
subForest c'
else
[c']
) ++ filter (/= c) children
| otherwise =
pointCell p v : children
where
contains f = isJust . maybecontain f
maybecontain f p = find (\x -> rectContains (rect x) (p,p)) f
divideUp :: [Int] -> [Float] -> Point -> [PKTree a] -> [[PKTree a]]
divideUp r w l children = npartition nodeBucket children
where
nodeBucket :: PKTree a -> Integer
nodeBucket x = sum $ zipWith4 bucket (upperPoint x) w l cumulDimSizes
cumulDimSizes = map product (inits (map fromIntegral r))
upperPoint = snd . rect
bucket x w l d = d * truncate ((xl)/w)
instantiateDivisions :: Int -> [Int] -> Rectangle -> [PKTree a] -> [PKTree a]
instantiateDivisions k r (l, u) children =
concat rest ++ concatMap (\div ->
let uDiv = nubBy (\x y -> rect x == rect y) div in
let newKids = instantiateDivisions k r (divBox uDiv) uDiv in
if length newKids < k then
newKids
else
[Node {
rootLabel = Inner (divBox uDiv),
subForest = newKids
}]
) subdivisions
where
divBox = box . fst . rect . head
box p = unzip $ zipWith (\x w ->
let low = w * fFloor (x/w) in
(low, low + w)
) p w
fFloor = fromIntegral . (floor :: Float -> Int)
(subdivisions, rest) =
partition (\x -> length x >= k) (divideUp r w l children)
w = zipWith3 (\l u r -> (ul) / fromIntegral r) l u r
radiusSearch :: Point
-> Float
-> PKTree a
-> [(Point, a)]
radiusSearch p r tree =
concatMap (\t -> case () of
_ | circleContains (rect t) -> getLeaves t
| circleIntersect (rect t) -> radiusSearch p r t
| otherwise -> []
) (subForest tree)
where
getLeaves (Node {rootLabel = (Leaf p v), subForest = _}) = [(p, v)]
getLeaves (Node {rootLabel = _, subForest = children}) =
concatMap getLeaves children
circleContains (l, u) =
sqdist (zip l p) <= sq r && sqdist (zip u p) <= sq r
circleIntersect rect =
let distWidth = zip (circleDistances rect) (halfRectDims rect) in
((not . or) (map (\(d,w) -> d > (w+r)) distWidth) &&
any (uncurry (<=)) distWidth) ||
sqdist distWidth <= sq r
sqdist = foldr (\(d,w) dist -> dist + sq (dw)) 0
sq = (^(2::Int))
circleDistances rect@(l, _) = zipWith3 (\l w p -> p l w) l (halfRectDims rect) p
halfRectDims (l, u) = zipWith (\l u -> (ul)/2) l u