diff options
| -rw-r--r-- | src/PPL/Sampling.hs | 19 |
1 files changed, 17 insertions, 2 deletions
diff --git a/src/PPL/Sampling.hs b/src/PPL/Sampling.hs index e674b2b..574025a 100644 --- a/src/PPL/Sampling.hs +++ b/src/PPL/Sampling.hs @@ -15,7 +15,7 @@ import Control.Arrow import Control.Monad import Control.Monad.IO.Class import Data.IORef -import Data.List (sort) +import Data.List (foldl') import Data.Vector qualified as V import Data.Vector.Hashtables qualified as H import Numeric.Log @@ -24,6 +24,21 @@ import Streaming.Prelude (Of, Stream, yield) import System.Random (StdGen) import System.Random qualified as R +data MinHeap a k = Node a k (MinHeap a k) (MinHeap a k) | Empty + deriving (Show) + +insert m k v = merge m (Node k v Empty Empty) + +merge Empty n = n +merge n Empty = n +merge n0@(Node k0 v0 l0 r0) n1@(Node k1 v1 l1 r1) = + if k0 < k1 + then Node k0 v0 (merge r0 n1) l0 + else Node k1 v1 r1 (merge l1 n0) + +toList Empty = [] +toList (Node k v l r) = v : toList (merge l r) + mh :: (MonadIO m) => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m () mh g p m = do let (g0, g1) = R.split g @@ -50,7 +65,7 @@ mh g p m = do m' <- H.clone m ks <- H.keys m let (rs :: [Double], qs :: [Double]) = (R.randoms *** R.randoms) (R.split g) - ks' = map snd $ sort $ zip rs $ V.toList ks + ks' = toList $ foldl' (\m -> uncurry (insert m)) Empty $ zip rs $ V.toList ks n = fromIntegral (V.length ks) void $ zipWithM_ (\k q -> H.insert m' k q) (take (1 + floor (p * fromIntegral n)) ks') qs newIORef (m', g0) |