New improved iterative search algorithm.

2 files changed, 69 insertions, 1 deletions

diff --git a/src/Data/MinMaxPSQ.hs b/src/Data/MinMaxPSQ.hs
new file mode 100644
index 00000000..96937604
--- /dev/null
+++ b/src/Data/MinMaxPSQ.hs
@@ -0,0 +1,64 @@
+{-# LANGUAGE BangPatterns #-}
+module Data.MinMaxPSQ where
+
+import Data.Ord
+import qualified Data.Wrapper.PSQ as PSQ
+         ;import Data.Wrapper.PSQ as PSQ hiding (insert, null, size)
+import Prelude hiding (null, take)
+
+data MinMaxPSQ k p = MinMaxPSQ !(PSQ k p) !(PSQ k (Down p))
+
+empty :: MinMaxPSQ k p
+empty = MinMaxPSQ PSQ.empty PSQ.empty
+
+null :: MinMaxPSQ k p -> Bool
+null (MinMaxPSQ nq xq) = PSQ.null nq
+
+size :: MinMaxPSQ k p -> Int
+size (MinMaxPSQ nq xq) = PSQ.size nq
+
+toList :: (Ord k, Ord p) => MinMaxPSQ k p -> [Binding k p]
+toList (MinMaxPSQ nq xq) = PSQ.toList nq
+
+fromList :: (Ord k, Ord p) => [Binding k p] -> MinMaxPSQ k p
+fromList kps = MinMaxPSQ (PSQ.fromList kps)
+                         (PSQ.fromList $ map (\(k :-> p) -> (k :-> Down p)) kps)
+
+findMin :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
+findMin (MinMaxPSQ nq xq) = PSQ.findMin nq
+
+findMax :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
+findMax (MinMaxPSQ nq xq) = fmap (\(k :-> Down p) -> k :-> p) $ PSQ.findMin xq
+
+insert :: (Ord k, Ord p) => k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
+insert k p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert k p        nq)
+                                         (PSQ.insert k (Down p) xq)
+
+delete :: (Ord k, Ord p) => k -> MinMaxPSQ k p -> MinMaxPSQ k p
+delete k (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.delete k nq) (PSQ.delete k xq)
+
+deleteMin :: (Ord k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMin (MinMaxPSQ nq xq) = case PSQ.minView nq of
+    Just (k :-> _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
+    Nothing             -> MinMaxPSQ nq xq
+
+deleteMax :: (Ord k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMax (MinMaxPSQ nq xq) = case PSQ.minView xq of
+    Just (k :-> _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
+    Nothing             -> MinMaxPSQ nq xq
+
+minView :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
+minView (MinMaxPSQ nq xq) = fmap (\(k :-> p, nq') -> (k :-> p, MinMaxPSQ nq' (PSQ.delete k xq)))
+                            $ PSQ.minView nq
+
+maxView :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
+maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> p, MinMaxPSQ (PSQ.delete k nq) xq'))
+                            $ PSQ.minView xq
+
+insertTake :: (Ord k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
+insertTake n k p q = take n $ insert k p q
+
+take :: (Ord k, Ord p) => Int -> MinMaxPSQ k p -> MinMaxPSQ k p
+take !n !q | (size q <= n) = q
+           | null q        = q
+           | otherwise     = take n $ deleteMax q
diff --git a/src/Data/Wrapper/PSQ.hs b/src/Data/Wrapper/PSQ.hs
index e8fa2d98..2c08011b 100644
--- a/src/Data/Wrapper/PSQ.hs
+++ b/src/Data/Wrapper/PSQ.hs
@@ -15,10 +15,14 @@ type Binding k p = (k,p,())
 pattern (:->) :: k -> p -> Binding k p
 pattern k :-> p <- (k,p,()) where k :-> p = (k,p,())
 
-key :: Binding k v -> k
+key :: Binding k p -> k
 key (k,p,v) = k
 {-# INLINE key #-}
 
+prio :: Binding k p -> p
+prio (k,p,v) = p
+{-# INLINE prio #-}
+
 insert :: (Ord k, Ord p) => k -> p -> PSQ k p -> PSQ k p
 insert k p q = OrdPSQ.insert k p () q
 {-# INLINE insert #-}