use skew heaps instead of sortingHEADmaster

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)