{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveFunctor #-}
{-# language DeriveGeneric #-}
{-# language LambdaCase #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiWayIf #-}
{-# options_ghc -Wno-unused-imports #-}
{-# options_ghc -Wno-unused-top-binds #-}
module Data.RPTree (
tree
, forest
, rpTreeCfg, RPTreeConfig(..)
, knn
, serialiseRPForest
, deserialiseRPForest
, recallWith
, leaves, levels, points, candidates
, treeStats, treeSize, leafSizes
, RPTreeStats
, Embed(..)
, RPTree, RPForest
, SVector, fromListSv, fromVectorSv
, DVector, fromListDv, fromVectorDv
, Inner(..), Scale(..)
, innerSS, innerSD, innerDD
, metricSSL2, metricSDL2
, scaleS, scaleD
, writeCsv
, writeDot
, BenchConfig(..), normalSparse2
, liftC
, randSeed
, dataSource
, datS, datD
, sparse, dense
, normal2, circle2d
) where
import Control.Monad (replicateM)
import Control.Monad.IO.Class (MonadIO(..))
import Data.Foldable (Foldable(..), maximumBy, minimumBy)
import Data.Functor.Identity (Identity(..))
import Data.List (partition, sortBy)
import Data.Monoid (Sum(..))
import Data.Ord (comparing)
import Data.Semigroup (Min(..))
import GHC.Generics (Generic)
import GHC.Word (Word64)
import Data.Sequence (Seq, (|>))
import qualified Data.Map as M (Map, fromList, toList, foldrWithKey, insert, insertWith, intersection)
import qualified Data.Set as S (Set, fromList, intersection, insert)
import Control.DeepSeq (NFData(..))
import qualified Data.IntPSQ as PQ (IntPSQ, findMin, minView, empty, insert, fromList, toList)
import Control.Monad.Trans.State (StateT(..), runStateT, evalStateT, State, runState, evalState, get, put)
import Control.Monad.Trans.Class (MonadTrans(..))
import qualified Data.Vector as V (Vector, replicateM, fromList)
import qualified Data.Vector.Generic as VG (Vector(..), unfoldrM, length, replicateM, (!), map, freeze, thaw, take, drop, unzip)
import qualified Data.Vector.Unboxed as VU (Vector, Unbox, fromList)
import qualified Data.Vector.Storable as VS (Vector)
import qualified Data.Vector.Algorithms.Merge as V (sortBy)
import Data.RPTree.Conduit (tree, forest, dataSource, liftC, rpTreeCfg, RPTreeConfig(..))
import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2, circle2d)
import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), Embed(..), leaves, levels, points, Inner(..), Scale(..), scaleS, scaleD, (/.), innerDD, innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, fromVectorSv, DVector(..), fromListDv, fromVectorDv, partitionAtMedian, Margin, getMargin, sortByVG, serialiseRPForest, deserialiseRPForest)
import Data.RPTree.Internal.Testing (BenchConfig(..), randSeed, datS, datD)
import Data.RPTree.Draw (writeDot, writeCsv)
knn :: (Ord p, Inner SVector v, VU.Unbox d, Real d) =>
(u d -> v d -> p)
-> Int
-> RPForest d (V.Vector (Embed u d x))
-> v d
-> V.Vector (p, Embed u d x)
knn :: (u d -> v d -> p)
-> Int
-> RPForest d (Vector (Embed u d x))
-> v d
-> Vector (p, Embed u d x)
knn u d -> v d -> p
distf Int
k RPForest d (Vector (Embed u d x))
tts v d
q = Int -> Vector (p, Embed u d x) -> Vector (p, Embed u d x)
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
VG.take Int
k (Vector (p, Embed u d x) -> Vector (p, Embed u d x))
-> Vector (p, Embed u d x) -> Vector (p, Embed u d x)
forall a b. (a -> b) -> a -> b
$ ((p, Embed u d x) -> p)
-> Vector (p, Embed u d x) -> Vector (p, Embed u d x)
forall (v :: * -> *) a b.
(Vector v a, Ord b) =>
(a -> b) -> v a -> v a
sortByVG (p, Embed u d x) -> p
forall a b. (a, b) -> a
fst Vector (p, Embed u d x)
cs
where
cs :: Vector (p, Embed u d x)
cs = (Embed u d x -> (p, Embed u d x))
-> Vector (Embed u d x) -> Vector (p, Embed u d x)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map (\Embed u d x
xe -> (Embed u d x -> u d
forall (v :: * -> *) e a. Embed v e a -> v e
eEmbed Embed u d x
xe u d -> v d -> p
`distf` v d
q, Embed u d x
xe)) (Vector (Embed u d x) -> Vector (p, Embed u d x))
-> Vector (Embed u d x) -> Vector (p, Embed u d x)
forall a b. (a -> b) -> a -> b
$ IntMap (Vector (Embed u d x)) -> Vector (Embed u d x)
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold (IntMap (Vector (Embed u d x)) -> Vector (Embed u d x))
-> IntMap (Vector (Embed u d x)) -> Vector (Embed u d x)
forall a b. (a -> b) -> a -> b
$ (RPTree d () (Vector (Embed u d x)) -> v d -> Vector (Embed u d x)
forall (v :: * -> *) d xs l.
(Inner SVector v, Unbox d, Ord d, Num d, Semigroup xs) =>
RPTree d l xs -> v d -> xs
`candidates` v d
q) (RPTree d () (Vector (Embed u d x)) -> Vector (Embed u d x))
-> RPForest d (Vector (Embed u d x))
-> IntMap (Vector (Embed u d x))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RPForest d (Vector (Embed u d x))
tts
knnPQ :: (Ord p, Inner SVector v, VU.Unbox d, RealFrac d) =>
(u d -> v d -> p)
-> Int
-> RPForest d (V.Vector (Embed u d x))
-> v d
-> V.Vector (p, Embed u d x)
knnPQ :: (u d -> v d -> p)
-> Int
-> RPForest d (Vector (Embed u d x))
-> v d
-> Vector (p, Embed u d x)
knnPQ u d -> v d -> p
distf Int
k RPForest d (Vector (Embed u d x))
tts v d
q = ((p, Embed u d x) -> p)
-> Vector (p, Embed u d x) -> Vector (p, Embed u d x)
forall (v :: * -> *) a b.
(Vector v a, Ord b) =>
(a -> b) -> v a -> v a
sortByVG (p, Embed u d x) -> p
forall a b. (a, b) -> a
fst Vector (p, Embed u d x)
cs
where
cs :: Vector (p, Embed u d x)
cs = (Embed u d x -> (p, Embed u d x))
-> Vector (Embed u d x) -> Vector (p, Embed u d x)
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map (\Embed u d x
xe -> (Embed u d x -> u d
forall (v :: * -> *) e a. Embed v e a -> v e
eEmbed Embed u d x
xe u d -> v d -> p
`distf` v d
q, Embed u d x
xe)) (Vector (Embed u d x) -> Vector (p, Embed u d x))
-> Vector (Embed u d x) -> Vector (p, Embed u d x)
forall a b. (a -> b) -> a -> b
$ IntMap (Vector (Embed u d x)) -> Vector (Embed u d x)
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold IntMap (Vector (Embed u d x))
cstt
cstt :: IntMap (Vector (Embed u d x))
cstt = (Int -> IntPSQ d (Vector (Embed u d x)) -> Vector (Embed u d x)
forall p (t :: * -> *) a.
(Ord p, Foldable t, Monoid (t a)) =>
Int -> IntPSQ p (t a) -> t a
takeFromPQ Int
nsing) (IntPSQ d (Vector (Embed u d x)) -> Vector (Embed u d x))
-> (RPTree d () (Vector (Embed u d x))
-> IntPSQ d (Vector (Embed u d x)))
-> RPTree d () (Vector (Embed u d x))
-> Vector (Embed u d x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RPTree d () (Vector (Embed u d x))
-> v d -> IntPSQ d (Vector (Embed u d x))
forall d (v :: * -> *) l xs.
(Fractional d, Ord d, Inner SVector v, Unbox d) =>
RPTree d l xs -> v d -> IntPSQ d xs
`candidatesPQ` v d
q) (RPTree d () (Vector (Embed u d x)) -> Vector (Embed u d x))
-> RPForest d (Vector (Embed u d x))
-> IntMap (Vector (Embed u d x))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RPForest d (Vector (Embed u d x))
tts
nsing :: Int
nsing = (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
n) Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
1
n :: Int
n = RPForest d (Vector (Embed u d x)) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length RPForest d (Vector (Embed u d x))
tts
recallWith :: (Inner SVector v, VU.Unbox d, Fractional b, Ord d, Ord a, Ord x, Ord (u d), Num d) =>
(u d -> v d -> a)
-> RPForest d (V.Vector (Embed u d x))
-> Int
-> v d
-> b
recallWith :: (u d -> v d -> a)
-> RPForest d (Vector (Embed u d x)) -> Int -> v d -> b
recallWith u d -> v d -> a
distf RPForest d (Vector (Embed u d x))
tt Int
k v d
q = IntMap b -> b
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum IntMap b
rs b -> b -> b
forall a. Fractional a => a -> a -> a
/ Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
where
rs :: IntMap b
rs = (RPTree d () (Vector (Embed u d x)) -> b)
-> RPForest d (Vector (Embed u d x)) -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\RPTree d () (Vector (Embed u d x))
t -> (u d -> v d -> a)
-> RPTree d () (Vector (Embed u d x)) -> Int -> v d -> b
forall (v :: * -> *) d p a x (u :: * -> *) l.
(Inner SVector v, Ord d, Unbox d, Fractional p, Ord a, Ord x,
Ord (u d), Num d) =>
(u d -> v d -> a)
-> RPTree d l (Vector (Embed u d x)) -> Int -> v d -> p
recallWith1 u d -> v d -> a
distf RPTree d () (Vector (Embed u d x))
t Int
k v d
q) RPForest d (Vector (Embed u d x))
tt
n :: Int
n = RPForest d (Vector (Embed u d x)) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length RPForest d (Vector (Embed u d x))
tt
recallWith1 :: (Inner SVector v, Ord d, VU.Unbox d, Fractional p, Ord a, Ord x, Ord (u d), Num d) =>
(u d -> v d -> a)
-> RPTree d l (V.Vector (Embed u d x))
-> Int
-> v d
-> p
recallWith1 :: (u d -> v d -> a)
-> RPTree d l (Vector (Embed u d x)) -> Int -> v d -> p
recallWith1 u d -> v d -> a
distf RPTree d l (Vector (Embed u d x))
tt Int
k v d
q = Int -> p
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Set (Embed u d x) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Set (Embed u d x)
aintk) p -> p -> p
forall a. Fractional a => a -> a -> a
/ Int -> p
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
where
aintk :: Set (Embed u d x)
aintk = Set (Embed u d x)
aa Set (Embed u d x) -> Set (Embed u d x) -> Set (Embed u d x)
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set (Embed u d x)
kk
aa :: Set (Embed u d x)
aa = Vector (Embed u d x) -> Set (Embed u d x)
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> Set a
set (Vector (Embed u d x) -> Set (Embed u d x))
-> Vector (Embed u d x) -> Set (Embed u d x)
forall a b. (a -> b) -> a -> b
$ RPTree d l (Vector (Embed u d x)) -> v d -> Vector (Embed u d x)
forall (v :: * -> *) d xs l.
(Inner SVector v, Unbox d, Ord d, Num d, Semigroup xs) =>
RPTree d l xs -> v d -> xs
candidates RPTree d l (Vector (Embed u d x))
tt v d
q
kk :: Set (Embed u d x)
kk = [Embed u d x] -> Set (Embed u d x)
forall a. Ord a => [a] -> Set a
S.fromList ([Embed u d x] -> Set (Embed u d x))
-> [Embed u d x] -> Set (Embed u d x)
forall a b. (a -> b) -> a -> b
$ ((Embed u d x, a) -> Embed u d x)
-> [(Embed u d x, a)] -> [Embed u d x]
forall a b. (a -> b) -> [a] -> [b]
map (Embed u d x, a) -> Embed u d x
forall a b. (a, b) -> a
fst ([(Embed u d x, a)] -> [Embed u d x])
-> [(Embed u d x, a)] -> [Embed u d x]
forall a b. (a -> b) -> a -> b
$ Int -> [(Embed u d x, a)] -> [(Embed u d x, a)]
forall a. Int -> [a] -> [a]
take Int
k [(Embed u d x, a)]
dists
dists :: [(Embed u d x, a)]
dists = ((Embed u d x, a) -> (Embed u d x, a) -> Ordering)
-> [(Embed u d x, a)] -> [(Embed u d x, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Embed u d x, a) -> a)
-> (Embed u d x, a) -> (Embed u d x, a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Embed u d x, a) -> a
forall a b. (a, b) -> b
snd) ([(Embed u d x, a)] -> [(Embed u d x, a)])
-> [(Embed u d x, a)] -> [(Embed u d x, a)]
forall a b. (a -> b) -> a -> b
$ Vector (Embed u d x, a) -> [(Embed u d x, a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Vector (Embed u d x, a) -> [(Embed u d x, a)])
-> Vector (Embed u d x, a) -> [(Embed u d x, a)]
forall a b. (a -> b) -> a -> b
$ (Embed u d x -> (Embed u d x, a))
-> Vector (Embed u d x) -> Vector (Embed u d x, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Embed u d x
x -> (Embed u d x
x, Embed u d x -> u d
forall (v :: * -> *) e a. Embed v e a -> v e
eEmbed Embed u d x
x u d -> v d -> a
`distf` v d
q)) Vector (Embed u d x)
xs
xs :: Vector (Embed u d x)
xs = RPTree d l (Vector (Embed u d x)) -> Vector (Embed u d x)
forall m d l. Monoid m => RPTree d l m -> m
points RPTree d l (Vector (Embed u d x))
tt
set :: (Foldable t, Ord a) => t a -> S.Set a
set :: t a -> Set a
set = (Set a -> a -> Set a) -> Set a -> t a -> Set a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> Set a -> Set a) -> Set a -> a -> Set a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert) Set a
forall a. Monoid a => a
mempty
{-# SCC candidates #-}
candidates :: (Inner SVector v, VU.Unbox d, Ord d, Num d, Semigroup xs) =>
RPTree d l xs
-> v d
-> xs
candidates :: RPTree d l xs -> v d -> xs
candidates (RPTree Vector (SVector d)
rvs RPT d l xs
tt) v d
x = Int -> RPT d l xs -> xs
forall a l. Semigroup a => Int -> RPT d l a -> a
go Int
0 RPT d l xs
tt
where
go :: Int -> RPT d l a -> a
go Int
_ (Tip l
_ a
xs) = a
xs
go Int
ixLev (Bin l
_ d
thr Margin d
margin RPT d l a
ltree RPT d l a
rtree) =
let
(d
mglo, d
mghi) = Margin d -> (d, d)
forall a. Margin a -> (a, a)
getMargin Margin d
margin
r :: SVector d
r = Vector (SVector d)
rvs Vector (SVector d) -> Int -> SVector d
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.! Int
ixLev
proj :: d
proj = SVector d
r SVector d -> v d -> d
forall (u :: * -> *) (v :: * -> *) a.
(Inner u v, Unbox a, Num a) =>
u a -> v a -> a
`inner` v d
x
i' :: Int
i' = Int -> Int
forall a. Enum a => a -> a
succ Int
ixLev
dl :: d
dl = d -> d
forall a. Num a => a -> a
abs (d
mglo d -> d -> d
forall a. Num a => a -> a -> a
- d
proj)
dr :: d
dr = d -> d
forall a. Num a => a -> a
abs (d
mghi d -> d -> d
forall a. Num a => a -> a -> a
- d
proj)
in
if | d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
thr Bool -> Bool -> Bool
&&
d
dl d -> d -> Bool
forall a. Ord a => a -> a -> Bool
> d
dr -> Int -> RPT d l a -> a
go Int
i' RPT d l a
ltree a -> a -> a
forall a. Semigroup a => a -> a -> a
<> Int -> RPT d l a -> a
go Int
i' RPT d l a
rtree
| d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
thr -> Int -> RPT d l a -> a
go Int
i' RPT d l a
ltree
| d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
> d
thr Bool -> Bool -> Bool
&&
d
dl d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
dr -> Int -> RPT d l a -> a
go Int
i' RPT d l a
ltree a -> a -> a
forall a. Semigroup a => a -> a -> a
<> Int -> RPT d l a -> a
go Int
i' RPT d l a
rtree
| Bool
otherwise -> Int -> RPT d l a -> a
go Int
i' RPT d l a
rtree
candidatesPQ :: (Fractional d, Ord d, Inner SVector v, VU.Unbox d) =>
RPTree d l xs
-> v d
-> PQ.IntPSQ d xs
candidatesPQ :: RPTree d l xs -> v d -> IntPSQ d xs
candidatesPQ (RPTree Vector (SVector d)
rvs RPT d l xs
tt) v d
x = S (IntPSQ d xs) -> IntPSQ d xs
forall a. S a -> a
evalS (S (IntPSQ d xs) -> IntPSQ d xs) -> S (IntPSQ d xs) -> IntPSQ d xs
forall a b. (a -> b) -> a -> b
$ Int -> RPT d l xs -> IntPSQ d xs -> d -> S (IntPSQ d xs)
forall l v.
Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
0 RPT d l xs
tt IntPSQ d xs
forall p v. IntPSQ p v
PQ.empty (d
1d -> d -> d
forall a. Fractional a => a -> a -> a
/d
0)
where
go :: Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
_ (Tip l
_ v
xs) IntPSQ d v
acc d
dprev =
d -> v -> IntPSQ d v -> StateT Int Identity (IntPSQ d v)
forall p v. Ord p => p -> v -> IntPSQ p v -> S (IntPSQ p v)
insPQ d
dprev v
xs IntPSQ d v
acc
go Int
ixLev (Bin l
_ d
thr Margin d
margin RPT d l v
ltree RPT d l v
rtree) IntPSQ d v
acc d
dprev = do
let
(d
mglo, d
mghi) = Margin d -> (d, d)
forall a. Margin a -> (a, a)
getMargin Margin d
margin
r :: SVector d
r = Vector (SVector d)
rvs Vector (SVector d) -> Int -> SVector d
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.! Int
ixLev
proj :: d
proj = SVector d
r SVector d -> v d -> d
forall (u :: * -> *) (v :: * -> *) a.
(Inner u v, Unbox a, Num a) =>
u a -> v a -> a
`inner` v d
x
i' :: Int
i' = Int -> Int
forall a. Enum a => a -> a
succ Int
ixLev
dl :: d
dl = d -> d
forall a. Num a => a -> a
abs (d
mglo d -> d -> d
forall a. Num a => a -> a -> a
- d
proj)
dr :: d
dr = d -> d
forall a. Num a => a -> a
abs (d
mghi d -> d -> d
forall a. Num a => a -> a -> a
- d
proj)
if | d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
thr Bool -> Bool -> Bool
&&
d
dl d -> d -> Bool
forall a. Ord a => a -> a -> Bool
> d
dr -> do
IntPSQ d v
ll <- Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
ltree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dl)
IntPSQ d v
lr <- Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
rtree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dr)
IntPSQ d v -> StateT Int Identity (IntPSQ d v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntPSQ d v -> StateT Int Identity (IntPSQ d v))
-> IntPSQ d v -> StateT Int Identity (IntPSQ d v)
forall a b. (a -> b) -> a -> b
$ [(Int, d, v)] -> IntPSQ d v
forall p v. Ord p => [(Int, p, v)] -> IntPSQ p v
PQ.fromList (IntPSQ d v -> [(Int, d, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ d v
ll [(Int, d, v)] -> [(Int, d, v)] -> [(Int, d, v)]
forall a. Semigroup a => a -> a -> a
<> IntPSQ d v -> [(Int, d, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ d v
lr)
| d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
thr -> Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
ltree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dl)
| d
proj d -> d -> Bool
forall a. Ord a => a -> a -> Bool
> d
thr Bool -> Bool -> Bool
&&
d
dl d -> d -> Bool
forall a. Ord a => a -> a -> Bool
< d
dr -> do
IntPSQ d v
ll <- Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
ltree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dl)
IntPSQ d v
lr <- Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
rtree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dr)
IntPSQ d v -> StateT Int Identity (IntPSQ d v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntPSQ d v -> StateT Int Identity (IntPSQ d v))
-> IntPSQ d v -> StateT Int Identity (IntPSQ d v)
forall a b. (a -> b) -> a -> b
$ [(Int, d, v)] -> IntPSQ d v
forall p v. Ord p => [(Int, p, v)] -> IntPSQ p v
PQ.fromList (IntPSQ d v -> [(Int, d, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ d v
ll [(Int, d, v)] -> [(Int, d, v)] -> [(Int, d, v)]
forall a. Semigroup a => a -> a -> a
<> IntPSQ d v -> [(Int, d, v)]
forall p v. IntPSQ p v -> [(Int, p, v)]
PQ.toList IntPSQ d v
lr)
| Bool
otherwise -> Int
-> RPT d l v -> IntPSQ d v -> d -> StateT Int Identity (IntPSQ d v)
go Int
i' RPT d l v
rtree IntPSQ d v
acc (d -> d -> d
forall a. Ord a => a -> a -> a
min d
dprev d
dr)
takeFromPQ :: (Ord p, Foldable t, Monoid (t a)) =>
Int
-> PQ.IntPSQ p (t a)
-> t a
takeFromPQ :: Int -> IntPSQ p (t a) -> t a
takeFromPQ Int
n IntPSQ p (t a)
pq = ((p, t a) -> t a) -> [(p, t a)] -> t a
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (p, t a) -> t a
forall a b. (a, b) -> b
snd ([(p, t a)] -> t a) -> [(p, t a)] -> t a
forall a b. (a -> b) -> a -> b
$ [(p, t a)] -> [(p, t a)]
forall a. [a] -> [a]
reverse ([(p, t a)] -> [(p, t a)]) -> [(p, t a)] -> [(p, t a)]
forall a b. (a -> b) -> a -> b
$ [(p, t a)] -> Int -> IntPSQ p (t a) -> [(p, t a)]
forall a (t :: * -> *) a.
(Ord a, Foldable t) =>
[(a, t a)] -> Int -> IntPSQ a (t a) -> [(a, t a)]
go [] Int
0 IntPSQ p (t a)
pq
where
go :: [(a, t a)] -> Int -> IntPSQ a (t a) -> [(a, t a)]
go [(a, t a)]
acc Int
nacc IntPSQ a (t a)
q = case IntPSQ a (t a) -> Maybe (Int, a, t a, IntPSQ a (t a))
forall p v. Ord p => IntPSQ p v -> Maybe (Int, p, v, IntPSQ p v)
PQ.minView IntPSQ a (t a)
q of
Maybe (Int, a, t a, IntPSQ a (t a))
Nothing -> [(a, t a)]
acc
Just (Int
_, a
p, t a
xs, IntPSQ a (t a)
pqRest) ->
let
nxs :: Int
nxs = t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
xs
nacc' :: Int
nacc' = Int
nacc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nxs
in if Int
nacc' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
then [(a, t a)] -> Int -> IntPSQ a (t a) -> [(a, t a)]
go ((a
p, t a
xs) (a, t a) -> [(a, t a)] -> [(a, t a)]
forall a. a -> [a] -> [a]
: [(a, t a)]
acc) Int
nacc' IntPSQ a (t a)
pqRest
else [(a, t a)]
acc
type S = State Int
evalS :: S a -> a
evalS :: S a -> a
evalS = (S a -> Int -> a) -> Int -> S a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip S a -> Int -> a
forall s a. State s a -> s -> a
evalState Int
0
insPQ :: (Ord p) => p -> v -> PQ.IntPSQ p v -> S (PQ.IntPSQ p v)
insPQ :: p -> v -> IntPSQ p v -> S (IntPSQ p v)
insPQ p
p v
x IntPSQ p v
pq = do
Int
i <- StateT Int Identity Int
forall (m :: * -> *) s. Monad m => StateT s m s
get
let
pq' :: IntPSQ p v
pq' = Int -> p -> v -> IntPSQ p v -> IntPSQ p v
forall p v. Ord p => Int -> p -> v -> IntPSQ p v -> IntPSQ p v
PQ.insert Int
i p
p v
x IntPSQ p v
pq
Int -> StateT Int Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (Int -> Int
forall a. Enum a => a -> a
succ Int
i)
IntPSQ p v -> S (IntPSQ p v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntPSQ p v
pq'
data RPTreeStats = RPTreeStats {
RPTreeStats -> Int
rptsLength :: Int
} deriving (RPTreeStats -> RPTreeStats -> Bool
(RPTreeStats -> RPTreeStats -> Bool)
-> (RPTreeStats -> RPTreeStats -> Bool) -> Eq RPTreeStats
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RPTreeStats -> RPTreeStats -> Bool
$c/= :: RPTreeStats -> RPTreeStats -> Bool
== :: RPTreeStats -> RPTreeStats -> Bool
$c== :: RPTreeStats -> RPTreeStats -> Bool
Eq, Int -> RPTreeStats -> ShowS
[RPTreeStats] -> ShowS
RPTreeStats -> String
(Int -> RPTreeStats -> ShowS)
-> (RPTreeStats -> String)
-> ([RPTreeStats] -> ShowS)
-> Show RPTreeStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RPTreeStats] -> ShowS
$cshowList :: [RPTreeStats] -> ShowS
show :: RPTreeStats -> String
$cshow :: RPTreeStats -> String
showsPrec :: Int -> RPTreeStats -> ShowS
$cshowsPrec :: Int -> RPTreeStats -> ShowS
Show)
treeStats :: RPTree d l a -> RPTreeStats
treeStats :: RPTree d l a -> RPTreeStats
treeStats (RPTree Vector (SVector d)
_ RPT d l a
tt) = Int -> RPTreeStats
RPTreeStats Int
l
where
l :: Int
l = RPT d l a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length RPT d l a
tt
treeSize :: (Foldable t) => RPTree d l (t a) -> Int
treeSize :: RPTree d l (t a) -> Int
treeSize = RPT d l Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (RPT d l Int -> Int)
-> (RPTree d l (t a) -> RPT d l Int) -> RPTree d l (t a) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RPTree d l (t a) -> RPT d l Int
forall (t :: * -> *) d l a.
Foldable t =>
RPTree d l (t a) -> RPT d l Int
leafSizes
leafSizes :: Foldable t => RPTree d l (t a) -> RPT d l Int
leafSizes :: RPTree d l (t a) -> RPT d l Int
leafSizes (RPTree Vector (SVector d)
_ RPT d l (t a)
tt) = t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (t a -> Int) -> RPT d l (t a) -> RPT d l Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RPT d l (t a)
tt