Add min/max handling operations for IntSet/IntMap

diff --git a/Data/IntMap.hs b/Data/IntMap.hs
index 4eaf99f..6c795aa 100644 (file)
--- a/Data/IntMap.hs
+++ b/Data/IntMap.hs
@@ -135,6 +135,23 @@ module Data.IntMap  (
             , isSubmapOf, isSubmapOfBy
             , isProperSubmapOf, isProperSubmapOfBy
             
+            -- * Min\/Max
+
+            , maxView
+            , minView
+            , findMin   
+            , findMax
+            , deleteMin
+            , deleteMax
+            , deleteFindMin
+            , deleteFindMax
+            , updateMin
+            , updateMax
+            , updateMinWithKey
+            , updateMaxWithKey 
+            , minViewWithKey
+            , maxViewWithKey
+
             -- * Debugging
             , showTree
             , showTreeWith
@@ -147,7 +164,7 @@ import qualified Data.IntSet as IntSet
 import Data.Monoid (Monoid(..))
 import Data.Typeable
 import Data.Foldable (Foldable(foldMap))
-
+import Control.Monad ( liftM )
 {-
 -- just for testing
 import qualified Prelude
@@ -669,6 +686,115 @@ intersectionWithKey f t Nil = Nil
 
 
 {--------------------------------------------------------------------
+  Min\/Max
+--------------------------------------------------------------------}
+
+-- | /O(log n)/. Update the value at the minimal key.
+updateMinWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
+updateMinWithKey f t
+    = case t of
+        Bin p m l r | m < 0 -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r
+        Bin p m l r         -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+        Nil -> error "maxView: empty map has no maximal element"
+
+updateMinWithKeyUnsigned f t
+    = case t of
+        Bin p m l r -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+
+-- | /O(log n)/. Update the value at the maximal key.
+updateMaxWithKey :: (Key -> a -> a) -> IntMap a -> IntMap a
+updateMaxWithKey f t
+    = case t of
+        Bin p m l r | m < 0 -> let t' = updateMaxWithKeyUnsigned f r in Bin p m r t'
+        Bin p m l r         -> let t' = updateMaxWithKeyUnsigned f l in Bin p m t' l
+        Tip k y -> Tip k (f k y)
+        Nil -> error "maxView: empty map has no maximal element"
+
+updateMaxWithKeyUnsigned f t
+    = case t of
+        Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t'
+        Tip k y -> Tip k (f k y)
+
+
+-- | /O(log n)/. Retrieves the maximal (key,value) couple of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+maxViewWithKey :: (Monad m) => IntMap a -> m ((Key, a), IntMap a)
+maxViewWithKey t
+    = case t of
+        Bin p m l r | m < 0 -> let (result, t') = maxViewUnsigned l in return (result, bin p m t' r)
+        Bin p m l r         -> let (result, t') = maxViewUnsigned r in return (result, bin p m l t')
+        Tip k y -> return ((k,y), Nil)
+        Nil -> fail "maxView: empty map has no maximal element"
+
+maxViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = maxViewUnsigned r in (result,bin p m l t')
+        Tip k y -> ((k,y), Nil)
+
+-- | /O(log n)/. Retrieves the minimal (key,value) couple of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+minViewWithKey :: (Monad m) => IntMap a -> m ((Key, a), IntMap a)
+minViewWithKey t
+    = case t of
+        Bin p m l r | m < 0 -> let (result, t') = minViewUnsigned r in return (result, bin p m l t')
+        Bin p m l r         -> let (result, t') = minViewUnsigned l in return (result, bin p m t' r)
+        Tip k y -> return ((k,y),Nil)
+        Nil -> fail "minView: empty map has no minimal element"
+
+minViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = minViewUnsigned l in (result,bin p m t' r)
+        Tip k y -> ((k,y),Nil)
+
+
+-- | /O(log n)/. Update the value at the maximal key.
+updateMax :: (a -> a) -> IntMap a -> IntMap a
+updateMax f = updateMaxWithKey (const f)
+
+-- | /O(log n)/. Update the value at the minimal key.
+updateMin :: (a -> a) -> IntMap a -> IntMap a
+updateMin f = updateMinWithKey (const f)
+
+
+-- Duplicate the Identity monad here because base < mtl.
+newtype Identity a = Identity { runIdentity :: a }
+instance Monad Identity where
+       return a = Identity a
+       m >>= k  = k (runIdentity m)
+-- Similar to the Arrow instance.
+first f (x,y) = (f x,y)
+
+
+-- | /O(log n)/. Retrieves the maximal key of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+maxView t = liftM (first snd) (maxViewWithKey t)
+
+-- | /O(log n)/. Retrieves the minimal key of the map, and the map stripped from that element.
+-- @fail@s (in the monad) when passed an empty map.
+minView t = liftM (first snd) (minViewWithKey t)
+
+-- | /O(log n)/. Delete and find the maximal element.
+deleteFindMax = runIdentity . maxView
+
+-- | /O(log n)/. Delete and find the minimal element.
+deleteFindMin = runIdentity . minView
+
+-- | /O(log n)/. The minimal key of the map.
+findMin = fst . runIdentity . minView
+
+-- | /O(log n)/. The maximal key of the map.
+findMax = fst . runIdentity . maxView
+
+-- | /O(log n)/. Delete the minimal key.
+deleteMin = snd . runIdentity . minView
+
+-- | /O(log n)/. Delete the maximal key.
+deleteMax = snd . runIdentity . maxView
+
+
+{--------------------------------------------------------------------
   Submap
 --------------------------------------------------------------------}
 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). 

diff --git a/Data/IntSet.hs b/Data/IntSet.hs
index be51ce7..1622608 100644 (file)
--- a/Data/IntSet.hs
+++ b/Data/IntSet.hs
@@ -68,6 +68,16 @@ module Data.IntSet  (
             , split
             , splitMember
 
+            -- * Min\/Max
+            , findMin   
+            , findMax
+            , deleteMin
+            , deleteMax
+            , deleteFindMin
+            , deleteFindMax
+            , maxView
+            , minView
+
             -- * Map
            , map
 
@@ -156,6 +166,8 @@ m1 \\ m2 = difference m1 m2
 data IntSet = Nil
             | Tip {-# UNPACK #-} !Int
             | Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !IntSet !IntSet
+-- Invariant: Nil is never found as a child of Bin.
+
 
 type Prefix = Int
 type Mask   = Int
@@ -209,7 +221,7 @@ member x t
       Tip y -> (x==y)
       Nil   -> False
     
--- | /O(log n)/. Is the element not in the set?
+-- | /O(min(n,W))/. Is the element not in the set?
 notMember :: Int -> IntSet -> Bool
 notMember k = not . member k
 
@@ -456,7 +468,7 @@ partition pred t
       Nil -> (Nil,Nil)
 
 
--- | /O(log n)/. The expression (@'split' x set@) is a pair @(set1,set2)@
+-- | /O(min(n,W))/. The expression (@'split' x set@) is a pair @(set1,set2)@
 -- where all elements in @set1@ are lower than @x@ and all elements in
 -- @set2@ larger than @x@.
 --
@@ -489,7 +501,7 @@ split' x t
         | otherwise -> (Nil,Nil)
       Nil -> (Nil,Nil)
 
--- | /O(log n)/. Performs a 'split' but also returns whether the pivot
+-- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot
 -- element was found in the original set.
 splitMember :: Int -> IntSet -> (IntSet,Bool,IntSet)
 splitMember x t
@@ -520,6 +532,80 @@ splitMember' x t
       Nil -> (Nil,False,Nil)
 
 {----------------------------------------------------------------------
+  Min/Max
+----------------------------------------------------------------------}
+
+-- | /O(min(n,W))/. Retrieves the maximal key of the set, and the set stripped from that element
+-- @fail@s (in the monad) when passed an empty set.
+maxView :: (Monad m) => IntSet -> m (Int, IntSet)
+maxView t
+    = case t of
+        Bin p m l r | m < 0 -> let (result,t') = maxViewUnsigned l in return (result, bin p m t' r)
+        Bin p m l r         -> let (result,t') = maxViewUnsigned r in return (result, bin p m l t')            
+        Tip y -> return (y,Nil)
+        Nil -> fail "maxView: empty set has no maximal element"
+
+maxViewUnsigned :: IntSet -> (Int, IntSet)
+maxViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = maxViewUnsigned r in (result, bin p m l t')
+        Tip y -> (y, Nil)
+
+-- | /O(min(n,W))/. Retrieves the minimal key of the set, and the set stripped from that element
+-- @fail@s (in the monad) when passed an empty set.
+minView :: (Monad m) => IntSet -> m (Int, IntSet)
+minView t
+    = case t of
+        Bin p m l r | m < 0 -> let (result,t') = minViewUnsigned r in return (result, bin p m l t')            
+        Bin p m l r         -> let (result,t') = minViewUnsigned l in return (result, bin p m t' r)
+        Tip y -> return (y, Nil)
+        Nil -> fail "minView: empty set has no minimal element"
+
+minViewUnsigned :: IntSet -> (Int, IntSet)
+minViewUnsigned t 
+    = case t of
+        Bin p m l r -> let (result,t') = minViewUnsigned l in (result, bin p m t' r)
+        Tip y -> (y, Nil)
+
+
+-- Duplicate the Identity monad here because base < mtl.
+newtype Identity a = Identity { runIdentity :: a }
+instance Monad Identity where
+       return a = Identity a
+       m >>= k  = k (runIdentity m)
+
+
+-- | /O(min(n,W))/. Delete and find the minimal element.
+-- 
+-- > deleteFindMin set = (findMin set, deleteMin set)
+deleteFindMin :: IntSet -> (Int, IntSet)
+deleteFindMin = runIdentity . minView
+
+-- | /O(min(n,W))/. Delete and find the maximal element.
+-- 
+-- > deleteFindMax set = (findMax set, deleteMax set)
+deleteFindMax :: IntSet -> (Int, IntSet)
+deleteFindMax = runIdentity . maxView
+
+-- | /O(min(n,W))/. The minimal element of a set.
+findMin :: IntSet -> Int
+findMin = fst . runIdentity . minView
+
+-- | /O(min(n,W))/. The maximal element of a set.
+findMax :: IntSet -> Int
+findMax = fst . runIdentity . maxView
+
+-- | /O(min(n,W))/. Delete the minimal element.
+deleteMin :: IntSet -> IntSet
+deleteMin = snd . runIdentity . minView
+
+-- | /O(min(n,W))/. Delete the maximal element.
+deleteMax :: IntSet -> IntSet
+deleteMax = snd . runIdentity . maxView
+
+
+
+{----------------------------------------------------------------------
   Map
 ----------------------------------------------------------------------}
 

diff --git a/Data/Map.hs b/Data/Map.hs
index 0710a28..399f74c 100644 (file)
--- a/Data/Map.hs
+++ b/Data/Map.hs
@@ -157,6 +157,8 @@ module Data.Map  (
             , updateMaxWithKey
             , minView
             , maxView
+            , minViewWithKey
+            , maxViewWithKey
             
             -- * Debugging
             , showTree