{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# options_ghc -Wno-unused-imports #-}
module Data.VPTree.Query (
range
, distances
) where
import Control.Monad.IO.Class (MonadIO(..))
import Data.Foldable (toList, foldrM, foldlM)
import Data.Sequence as SQ (Seq)
import Data.Sequence ((|>))
import Control.Monad.State (MonadState(..))
import qualified Data.IntPSQ as PQ (IntPSQ, insert, size, empty, toList, minView)
import Control.Monad.Trans.State (State, evalState, runState)
import qualified Data.Vector as V (Vector)
import Data.VPTree.Internal (VT(..), VPTree(..))
psqList :: (Ord p) =>
PQ.IntPSQ p b -> [(p, b)]
psqList :: IntPSQ p b -> [(p, b)]
psqList IntPSQ p b
q = case IntPSQ p b -> Maybe (Int, p, b, IntPSQ p b)
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v, IntPSQ p v)
PQ.minView IntPSQ p b
q of
Maybe (Int, p, b, IntPSQ p b)
Nothing -> [(p, b)]
forall a. Monoid a => a
mempty
Just (Int
_, p
p, b
v, IntPSQ p b
qrest) -> (p
p, b
v) (p, b) -> [(p, b)] -> [(p, b)]
forall a. a -> [a] -> [a]
: IntPSQ p b -> [(p, b)]
forall p b. Ord p => IntPSQ p b -> [(p, b)]
psqList IntPSQ p b
qrest
distances :: VPTree b a
-> a
-> [b]
distances :: VPTree b a -> a -> [b]
distances (VPT VT b a
tt a -> a -> b
distf) a
x = (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> b
distf a
x) ([a] -> [b]) -> [a] -> [b]
forall a b. (a -> b) -> a -> b
$ VT b a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList VT b a
tt
range :: (Num p, Ord p) =>
VPTree p a
-> p
-> a
-> [(p, a)]
range :: VPTree p a -> p -> a -> [(p, a)]
range (VPT VT p a
tt a -> a -> p
distf) p
eps a
x = IntPSQ p a -> [(p, a)]
forall p b. Ord p => IntPSQ p b -> [(p, b)]
psqList (IntPSQ p a -> [(p, a)]) -> IntPSQ p a -> [(p, a)]
forall a b. (a -> b) -> a -> b
$ p -> a -> (a -> a -> p) -> VT p a -> IntPSQ p a
forall b a.
(Num b, Ord b) =>
b -> a -> (a -> a -> b) -> VT b a -> IntPSQ b a
rangeVT p
eps a
x a -> a -> p
distf VT p a
tt
rangeVT :: (Num b, Ord b) =>
b
-> a -> (a -> a -> b) -> VT b a -> PQ.IntPSQ b a
rangeVT :: b -> a -> (a -> a -> b) -> VT b a -> IntPSQ b a
rangeVT b
eps a
x a -> a -> b
distf = (State Int (IntPSQ b a) -> Int -> IntPSQ b a)
-> Int -> State Int (IntPSQ b a) -> IntPSQ b a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State Int (IntPSQ b a) -> Int -> IntPSQ b a
forall s a. State s a -> s -> a
evalState Int
0 (State Int (IntPSQ b a) -> IntPSQ b a)
-> (VT b a -> State Int (IntPSQ b a)) -> VT b a -> IntPSQ b a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntPSQ b a -> VT b a -> State Int (IntPSQ b a)
forall (m :: * -> *).
MonadState Int m =>
IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
forall p v. IntPSQ p v
PQ.empty
where
go :: IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc = \case
Tip Vector a
ts ->
(IntPSQ b a -> a -> m (IntPSQ b a))
-> IntPSQ b a -> Vector a -> m (IntPSQ b a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM IntPSQ b a -> a -> m (IntPSQ b a)
forall (m :: * -> *).
MonadState Int m =>
IntPSQ b a -> a -> m (IntPSQ b a)
insf IntPSQ b a
acc Vector a
ts
where
insf :: IntPSQ b a -> a -> m (IntPSQ b a)
insf IntPSQ b a
ac a
t
| b
d b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
eps = do
Int
i <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
let ac' :: IntPSQ b a
ac' = Int -> b -> a -> IntPSQ b a -> IntPSQ b a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
PQ.insert Int
i b
d a
t IntPSQ b a
ac
Int -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
IntPSQ b a -> m (IntPSQ b a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntPSQ b a
ac'
| Bool
otherwise = IntPSQ b a -> m (IntPSQ b a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntPSQ b a
ac
where
d :: b
d = a -> a -> b
distf a
x a
t
Bin b
mu a
v VT b a
ll VT b a
rr
| b
d b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
eps -> do
Int
i <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
let acc' :: IntPSQ b a
acc' = Int -> b -> a -> IntPSQ b a -> IntPSQ b a
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
PQ.insert Int
i b
d a
v IntPSQ b a
acc
Int -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc' VT b a
ll
| b
d b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
mu b -> b -> b
forall a. Num a => a -> a -> a
+ b
eps -> IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc VT b a
rr
| b
d b -> b -> Bool
forall a. Ord a => a -> a -> Bool
<= b
mu b -> b -> b
forall a. Num a => a -> a -> a
+ b
eps Bool -> Bool -> Bool
&& b
d b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
mu b -> b -> b
forall a. Num a => a -> a -> a
- b
eps -> do
IntPSQ b a
accl <- IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc VT b a
ll
IntPSQ b a
accr <- IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc VT b a
rr
IntPSQ b a -> IntPSQ b a -> m (IntPSQ b a)
forall (m :: * -> *) b c.
(MonadState Int m, Ord b) =>
IntPSQ b c -> IntPSQ b c -> m (IntPSQ b c)
union IntPSQ b a
accl IntPSQ b a
accr
| Bool
otherwise -> IntPSQ b a -> VT b a -> m (IntPSQ b a)
go IntPSQ b a
acc VT b a
ll
where
d :: b
d = a -> a -> b
distf a
x a
v
union :: (MonadState Int m, Ord b) =>
PQ.IntPSQ b c -> PQ.IntPSQ b c -> m (PQ.IntPSQ b c)
union :: IntPSQ b c -> IntPSQ b c -> m (IntPSQ b c)
union IntPSQ b c
q1 IntPSQ b c
q2 = do
Int
i0 <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
IntPSQ b c -> m (IntPSQ b c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntPSQ b c -> m (IntPSQ b c)) -> IntPSQ b c -> m (IntPSQ b c)
forall a b. (a -> b) -> a -> b
$ (State Int (IntPSQ b c) -> Int -> IntPSQ b c)
-> Int -> State Int (IntPSQ b c) -> IntPSQ b c
forall a b c. (a -> b -> c) -> b -> a -> c
flip State Int (IntPSQ b c) -> Int -> IntPSQ b c
forall s a. State s a -> s -> a
evalState Int
i0 (State Int (IntPSQ b c) -> IntPSQ b c)
-> State Int (IntPSQ b c) -> IntPSQ b c
forall a b. (a -> b) -> a -> b
$ ((Int, b, c) -> IntPSQ b c -> State Int (IntPSQ b c))
-> IntPSQ b c -> [(Int, b, c)] -> State Int (IntPSQ b c)
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM (Int, b, c) -> IntPSQ b c -> State Int (IntPSQ b c)
forall (m :: * -> *) b a c.
(MonadState Int m, Ord b) =>
(a, b, c) -> IntPSQ b c -> m (IntPSQ b c)
f IntPSQ b c
forall p v. IntPSQ p v
PQ.empty ([(Int, b, c)] -> State Int (IntPSQ b c))
-> [(Int, b, c)] -> State Int (IntPSQ b c)
forall a b. (a -> b) -> a -> b
$ [(Int, b, c)]
l1 [(Int, b, c)] -> [(Int, b, c)] -> [(Int, b, c)]
forall a. Semigroup a => a -> a -> a
<> [(Int, b, c)]
l2
where
f :: (a, b, c) -> IntPSQ b c -> m (IntPSQ b c)
f (a
_, b
p, c
v) IntPSQ b c
acc = do
Int
i <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
let acc' :: IntPSQ b c
acc' = Int -> b -> c -> IntPSQ b c -> IntPSQ b c
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
PQ.insert Int
i b
p c
v IntPSQ b c
acc
Int -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Enum a => a -> a
succ Int
i
IntPSQ b c -> m (IntPSQ b c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntPSQ b c
acc'
l1 :: [(Int, b, c)]
l1 = IntPSQ b c -> [(Int, b, c)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ b c
q1
l2 :: [(Int, b, c)]
l2 = IntPSQ b c -> [(Int, b, c)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ b c
q2