O(1) size operation for MinMaxPSQ.

1 files changed, 41 insertions, 28 deletions

diff --git a/src/Data/MinMaxPSQ.hs b/src/Data/MinMaxPSQ.hs
index 3b9a4d6c..e7d7c760 100644
--- a/src/Data/MinMaxPSQ.hs
+++ b/src/Data/MinMaxPSQ.hs
@@ -10,72 +10,85 @@ import qualified Data.Wrapper.PSQ as PSQ
          ;import Data.Wrapper.PSQ as PSQ hiding (insert, insert', null, size)
 import Prelude hiding (null, take)
 
-data MinMaxPSQ' k p v = MinMaxPSQ !(PSQ' k p v) !(PSQ' k (Down p) v)
+data MinMaxPSQ' k p v = MinMaxPSQ !Int !(PSQ' k p v) !(PSQ' k (Down p) v)
 type MinMaxPSQ k p = MinMaxPSQ' k p ()
 
 empty :: MinMaxPSQ' k p v
-empty = MinMaxPSQ PSQ.empty PSQ.empty
+empty = MinMaxPSQ 0 PSQ.empty PSQ.empty
 
 singleton' :: (PSQKey k, Ord p) => k -> v -> p -> MinMaxPSQ' k p v
-singleton' k v p = MinMaxPSQ (PSQ.singleton' k v p) (PSQ.singleton' k v (Down p))
+singleton' k v p = MinMaxPSQ 1 (PSQ.singleton' k v p) (PSQ.singleton' k v (Down p))
 
 null :: MinMaxPSQ' k p v -> Bool
-null (MinMaxPSQ nq xq) = PSQ.null nq
+null (MinMaxPSQ sz _ _) = sz==0
+{-# INLINE null #-}
 
 size :: MinMaxPSQ' k p v -> Int
-size (MinMaxPSQ nq xq) = PSQ.size nq
+size (MinMaxPSQ sz _ _) = sz
+{-# INLINE size #-}
 
 toList :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> [Binding' k p v]
-toList (MinMaxPSQ nq xq) = PSQ.toList nq
+toList (MinMaxPSQ _ nq xq) = PSQ.toList nq
 
 fromList :: (PSQKey k, Ord p) => [Binding' k p v] -> MinMaxPSQ' k p v
-fromList kps = MinMaxPSQ (PSQ.fromList kps)
-                         (PSQ.fromList $ map (\(Binding k v p) -> Binding k v (Down p)) kps)
+fromList kps = let nq = PSQ.fromList kps
+                   xq = PSQ.fromList $ map (\(Binding k v p) -> Binding k v (Down p)) kps
+               in MinMaxPSQ (PSQ.size nq) nq xq
 
 findMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
-findMin (MinMaxPSQ nq xq) = PSQ.findMin nq
+findMin (MinMaxPSQ _ nq xq) = PSQ.findMin nq
 
 findMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v)
-findMax (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p)) -> Binding k v p) $ PSQ.findMin xq
+findMax (MinMaxPSQ _ nq xq) = fmap (\(Binding k v (Down p)) -> Binding k v p) $ PSQ.findMin xq
 
 insert :: (PSQKey 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)
+insert k p (MinMaxPSQ sz nq xq) = case PSQ.insertView k p () nq of
+    (Just _ ,nq') -> MinMaxPSQ sz     nq' (PSQ.insert k (Down p) xq)
+    (Nothing,nq') -> MinMaxPSQ (sz+1) nq' (PSQ.insert k (Down p) xq)
 
 insert' :: (PSQKey k, Ord p) => k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
-insert' k v p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert' k v p        nq)
-                                            (PSQ.insert' k v (Down p) xq)
+insert' k v p (MinMaxPSQ sz nq xq) = case PSQ.insertView k p v nq of
+    (Just _ ,nq') -> MinMaxPSQ sz     nq' (PSQ.insert' k v (Down p) xq)
+    (Nothing,nq') -> MinMaxPSQ (sz+1) nq' (PSQ.insert' k v (Down p) xq)
 
 delete :: (PSQKey k, Ord p) => k -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
-delete k (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.delete k nq) (PSQ.delete k xq)
+delete k q@(MinMaxPSQ sz nq xq) = case PSQ.deleteView k nq of
+    Just (_,_,nq') -> MinMaxPSQ (sz - 1) nq' (PSQ.delete k xq)
+    Nothing        -> q
 
 deleteMin :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
-deleteMin (MinMaxPSQ nq xq) = case PSQ.minView nq of
-    Just (Binding k _ _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
-    Nothing                   -> MinMaxPSQ nq xq
+deleteMin q@(MinMaxPSQ sz nq xq) = case PSQ.minView nq of
+    Just (Binding k _ _, nq') -> MinMaxPSQ (sz - 1) nq' (PSQ.delete k xq)
+    Nothing                   -> q
 
 deleteMax :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> MinMaxPSQ' k p v
-deleteMax (MinMaxPSQ nq xq) = case PSQ.minView xq of
-    Just (Binding k _ _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
-    Nothing                   -> MinMaxPSQ nq xq
+deleteMax q@(MinMaxPSQ sz nq xq) = case PSQ.minView xq of
+    Just (Binding k _ _, xq') -> MinMaxPSQ (sz - 1) (PSQ.delete k nq) xq'
+    Nothing                   -> q
 
 minView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
-minView (MinMaxPSQ nq xq) = fmap (\(Binding k v p, nq') -> (Binding k v p, MinMaxPSQ nq' (PSQ.delete k xq)))
-                            $ PSQ.minView nq
+minView (MinMaxPSQ sz nq xq) = fmap (\(Binding k v p, nq') -> (Binding k v p, MinMaxPSQ (sz-1) nq' (PSQ.delete k xq)))
+                               $ PSQ.minView nq
 
 maxView :: (PSQKey k, Ord p) => MinMaxPSQ' k p v -> Maybe (Binding' k p v, MinMaxPSQ' k p v)
-maxView (MinMaxPSQ nq xq) = fmap (\(Binding k v (Down p), xq') -> (Binding k v p, MinMaxPSQ (PSQ.delete k nq) xq'))
-                            $ PSQ.minView xq
+maxView (MinMaxPSQ sz nq xq) = fmap (\(Binding k v (Down p), xq') -> (Binding k v p, MinMaxPSQ (sz-1) (PSQ.delete k nq) xq'))
+                               $ PSQ.minView xq
 
 -- | Maintains a bounded 'MinMaxPSQ' by deleting the maximum element if the
 -- insertion would cause the queue to have too many elements.
 insertTake :: (PSQKey k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
-insertTake n k p q = take n $ insert k p q
+insertTake n k p q
+    | size q <  n  = insert k p q
+    | size q == n  = insert k p $ deleteMax q
+    | otherwise    = take n $ insert k p q
 
 -- | Maintains a bounded 'MinMaxPSQ\'' by deleting the maximum element if the
 -- insertion would cause the queue to have too many elements.
 insertTake' :: (PSQKey k, Ord p) => Int -> k -> v -> p -> MinMaxPSQ' k p v -> MinMaxPSQ' k p v
-insertTake' n k v p q = take n $ insert' k v p q
+insertTake' n k v p q
+    | size q <  n  = insert' k v p q
+    | size q == n  = insert' k v p $ deleteMax q
+    | otherwise    = take n $ insert' k v p q
 
 
 -- | Truncate the 'MinMaxPSQ' to the given number of lowest priority elements.
@@ -96,4 +109,4 @@ takeView !n !q | (size q <= n) = ([], q)
 
 
 lookup' :: (PSQKey k, Ord p) => k -> MinMaxPSQ' k p v -> Maybe (p, v)
-lookup' k (MinMaxPSQ q _) = PSQ.lookup k q
+lookup' k (MinMaxPSQ _ q _) = PSQ.lookup k q