Add min/max handling operations for IntSet/IntMap
[haskell-directory.git] / Data / Map.hs
index beddb7b..399f74c 100644 (file)
@@ -1,3 +1,5 @@
+{-# OPTIONS_GHC -fno-bang-patterns #-}
+
 -----------------------------------------------------------------------------
 -- |
 -- Module      :  Data.Map
 --
 -- An efficient implementation of maps from keys to values (dictionaries).
 --
--- This module is intended to be imported @qualified@, to avoid name
--- clashes with Prelude functions.  eg.
+-- Since many function names (but not the type name) clash with
+-- "Prelude" names, this module is usually imported @qualified@, e.g.
 --
--- >  import Data.Map as Map
+-- >  import Data.Map (Map)
+-- >  import qualified Data.Map as Map
 --
 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
 -- trees of /bounded balance/) as described by:
@@ -42,6 +45,7 @@ module Data.Map  (
             , null
             , size
             , member
+            , notMember
             , lookup
             , findWithDefault
             
@@ -52,6 +56,7 @@ module Data.Map  (
             -- ** Insertion
             , insert
             , insertWith, insertWithKey, insertLookupWithKey
+            , insertWith', insertWithKey'
             
             -- ** Delete\/Update
             , delete
@@ -60,6 +65,7 @@ module Data.Map  (
             , update
             , updateWithKey
             , updateLookupWithKey
+            , alter
 
             -- * Combine
 
@@ -119,6 +125,11 @@ module Data.Map  (
             , partition
             , partitionWithKey
 
+            , mapMaybe
+            , mapMaybeWithKey
+            , mapEither
+            , mapEitherWithKey
+
             , split         
             , splitLookup   
 
@@ -144,6 +155,10 @@ module Data.Map  (
             , updateMax
             , updateMinWithKey
             , updateMaxWithKey
+            , minView
+            , maxView
+            , minViewWithKey
+            , maxViewWithKey
             
             -- * Debugging
             , showTree
@@ -156,7 +171,7 @@ import qualified Data.Set as Set
 import qualified Data.List as List
 import Data.Monoid (Monoid(..))
 import Data.Typeable
-import Control.Applicative (Applicative(..))
+import Control.Applicative (Applicative(..), (<$>))
 import Data.Traversable (Traversable(traverse))
 import Data.Foldable (Foldable(foldMap))
 
@@ -216,7 +231,7 @@ instance (Data k, Data a, Ord k) => Data (Map k a) where
   toConstr _     = error "toConstr"
   gunfold _ _    = error "gunfold"
   dataTypeOf _   = mkNorepType "Data.Map.Map"
-  dataCast2      = gcast2
+  dataCast2 f    = gcast2 f
 
 #endif
 
@@ -238,7 +253,12 @@ size t
       Bin sz k x l r  -> sz
 
 
--- | /O(log n)/. Lookup the value at a key in the map.
+-- | /O(log n)/. Lookup the value at a key in the map. 
+--
+-- The function will 
+-- @return@ the result in the monad or @fail@ in it the key isn't in the 
+-- map. Often, the monad to use is 'Maybe', so you get either 
+-- @('Just' result)@ or @'Nothing'@.
 lookup :: (Monad m,Ord k) => k -> Map k a -> m a
 lookup k t = case lookup' k t of
     Just x -> return x
@@ -270,6 +290,10 @@ member k m
       Nothing -> False
       Just x  -> True
 
+-- | /O(log n)/. Is the key not a member of the map?
+notMember :: Ord k => k -> Map k a -> Bool
+notMember k m = not $ member k m
+
 -- | /O(log n)/. Find the value at a key.
 -- Calls 'error' when the element can not be found.
 find :: Ord k => k -> Map k a -> a
@@ -327,6 +351,12 @@ insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
 insertWith f k x m          
   = insertWithKey (\k x y -> f x y) k x m
 
+-- | Same as 'insertWith', but the combining function is applied strictly.
+insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
+insertWith' f k x m          
+  = insertWithKey' (\k x y -> f x y) k x m
+
+
 -- | /O(log n)/. Insert with a combining function.
 -- @'insertWithKey' f key value mp@ 
 -- will insert the pair (key, value) into @mp@ if key does
@@ -343,6 +373,18 @@ insertWithKey f kx x t
                GT -> balance ky y l (insertWithKey f kx x r)
                EQ -> Bin sy kx (f kx x y) l r
 
+-- | Same as 'insertWithKey', but the combining function is applied strictly.
+insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
+insertWithKey' f kx x t
+  = case t of
+      Tip -> singleton kx x
+      Bin sy ky y l r
+          -> case compare kx ky of
+               LT -> balance ky y (insertWithKey' f kx x l) r
+               GT -> balance ky y l (insertWithKey' f kx x r)
+               EQ -> let x' = f kx x y in seq x' (Bin sy kx x' l r)
+
+
 -- | /O(log n)/. The expression (@'insertLookupWithKey' f k x map@)
 -- is a pair where the first element is equal to (@'lookup' k map@)
 -- and the second element equal to (@'insertWithKey' f k x map@).
@@ -420,6 +462,23 @@ updateLookupWithKey f k t
                        Just x' -> (Just x',Bin sx kx x' l r)
                        Nothing -> (Just x,glue l r)
 
+-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
+-- 'alter' can be used to insert, delete, or update a value in a 'Map'.
+-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
+alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
+alter f k t
+  = case t of
+      Tip -> case f Nothing of
+               Nothing -> Tip
+               Just x -> singleton k x
+      Bin sx kx x l r 
+          -> case compare k kx of
+               LT -> balance kx x (alter f k l) r
+               GT -> balance kx x l (alter f k r)
+               EQ -> case f (Just x) of
+                       Just x' -> Bin sx kx x' l r
+                       Nothing -> glue l r
+
 {--------------------------------------------------------------------
   Indexing
 --------------------------------------------------------------------}
@@ -486,13 +545,13 @@ deleteAt i map
 findMin :: Map k a -> (k,a)
 findMin (Bin _ kx x Tip r)  = (kx,x)
 findMin (Bin _ kx x l r)    = findMin l
-findMin Tip                 = error "Map.findMin: empty tree has no minimal element"
+findMin Tip                 = error "Map.findMin: empty map has no minimal element"
 
 -- | /O(log n)/. The maximal key of the map.
 findMax :: Map k a -> (k,a)
 findMax (Bin _ kx x l Tip)  = (kx,x)
 findMax (Bin _ kx x l r)    = findMax r
-findMax Tip                 = error "Map.findMax: empty tree has no maximal element"
+findMax Tip                 = error "Map.findMax: empty map has no maximal element"
 
 -- | /O(log n)/. Delete the minimal key.
 deleteMin :: Map k a -> Map k a
@@ -537,6 +596,33 @@ updateMaxWithKey f t
       Bin sx kx x l r    -> balance kx x l (updateMaxWithKey f r)
       Tip                -> Tip
 
+-- | /O(log n)/. Retrieves the minimal (key,value) pair of the map, and the map stripped from that element
+-- @fail@s (in the monad) when passed an empty map.
+minViewWithKey :: Monad m => Map k a -> m ((k,a), Map k a)
+minViewWithKey Tip = fail "Map.minView: empty map"
+minViewWithKey x = return (deleteFindMin x)
+
+-- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and the map stripped from that element
+-- @fail@s (in the monad) when passed an empty map.
+maxViewWithKey :: Monad m => Map k a -> m ((k,a), Map k a)
+maxViewWithKey Tip = fail "Map.maxView: empty map"
+maxViewWithKey x = return (deleteFindMax x)
+
+-- | /O(log n)/. Retrieves the minimal key\'s value of the map, and the map stripped from that element
+-- @fail@s (in the monad) when passed an empty map.
+minView :: Monad m => Map k a -> m (a, Map k a)
+minView Tip = fail "Map.minView: empty map"
+minView x = return (first snd $ deleteFindMin x)
+
+-- | /O(log n)/. Retrieves the maximal key\'s value of the map, and the map stripped from that element
+-- @fail@s (in the monad) when passed an empty map.
+maxView :: Monad m => Map k a -> m (a, Map k a)
+maxView Tip = fail "Map.maxView: empty map"
+maxView x = return (first snd $ deleteFindMax x)
+
+-- Update the 1st component of a tuple (special case of Control.Arrow.first)
+first :: (a -> b) -> (a,c) -> (b,c)
+first f (x,y) = (f x, y)
 
 {--------------------------------------------------------------------
   Union. 
@@ -830,6 +916,33 @@ partitionWithKey p (Bin _ kx x l r)
     (l1,l2) = partitionWithKey p l
     (r1,r2) = partitionWithKey p r
 
+-- | /O(n)/. Map values and collect the 'Just' results.
+mapMaybe :: Ord k => (a -> Maybe b) -> Map k a -> Map k b
+mapMaybe f m
+  = mapMaybeWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and collect the 'Just' results.
+mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> Map k a -> Map k b
+mapMaybeWithKey f Tip = Tip
+mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
+  Just y  -> join kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+  Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+
+-- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
+mapEither :: Ord k => (a -> Either b c) -> Map k a -> (Map k b, Map k c)
+mapEither f m
+  = mapEitherWithKey (\k x -> f x) m
+
+-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
+mapEitherWithKey :: Ord k =>
+  (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c)
+mapEitherWithKey f Tip = (Tip, Tip)
+mapEitherWithKey f (Bin _ kx x l r) = case f kx x of
+  Left y  -> (join kx y l1 r1, merge l2 r2)
+  Right z -> (merge l1 r1, join kx z l2 r2)
+  where
+    (l1,l2) = mapEitherWithKey f l
+    (r1,r2) = mapEitherWithKey f r
 
 {--------------------------------------------------------------------
   Mapping