X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=Data%2FMap.hs;h=399f74c7abf4ffa67cefb265e47fa5ab876d86d1;hb=74bc2d04fdbae494bcf4839c4ec5e6ec1d0bf600;hp=b92bc36e116e2c9adae0ca9d6fc17353291970f7;hpb=d27ff3035d8276d5ad43e2570af3357d9bd05071;p=haskell-directory.git diff --git a/Data/Map.hs b/Data/Map.hs index b92bc36..399f74c 100644 --- a/Data/Map.hs +++ b/Data/Map.hs @@ -1,3 +1,5 @@ +{-# OPTIONS_GHC -fno-bang-patterns #-} + ----------------------------------------------------------------------------- -- | -- Module : Data.Map @@ -9,10 +11,11 @@ -- -- 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: @@ -24,11 +27,15 @@ -- * J. Nievergelt and E.M. Reingold, -- \"/Binary search trees of bounded balance/\", -- SIAM journal of computing 2(1), March 1973. +-- +-- Note that the implementation is /left-biased/ -- the elements of a +-- first argument are always preferred to the second, for example in +-- 'union' or 'insert'. ----------------------------------------------------------------------------- module Data.Map ( -- * Map type - Map -- instance Eq,Show + Map -- instance Eq,Show,Read -- * Operators , (!), (\\) @@ -38,6 +45,7 @@ module Data.Map ( , null , size , member + , notMember , lookup , findWithDefault @@ -48,6 +56,7 @@ module Data.Map ( -- ** Insertion , insert , insertWith, insertWithKey, insertLookupWithKey + , insertWith', insertWithKey' -- ** Delete\/Update , delete @@ -56,6 +65,7 @@ module Data.Map ( , update , updateWithKey , updateLookupWithKey + , alter -- * Combine @@ -115,6 +125,11 @@ module Data.Map ( , partition , partitionWithKey + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + , split , splitLookup @@ -140,6 +155,10 @@ module Data.Map ( , updateMax , updateMinWithKey , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey -- * Debugging , showTree @@ -148,10 +167,13 @@ module Data.Map ( ) where import Prelude hiding (lookup,map,filter,foldr,foldl,null) -import Data.Monoid import qualified Data.Set as Set import qualified Data.List as List +import Data.Monoid (Monoid(..)) import Data.Typeable +import Control.Applicative (Applicative(..), (<$>)) +import Data.Traversable (Traversable(traverse)) +import Data.Foldable (Foldable(foldMap)) {- -- for quick check @@ -162,6 +184,7 @@ import List(nub,sort) -} #if __GLASGOW_HASKELL__ +import Text.Read import Data.Generics.Basics import Data.Generics.Instances #endif @@ -189,6 +212,11 @@ data Map k a = Tip type Size = Int +instance (Ord k) => Monoid (Map k v) where + mempty = empty + mappend = union + mconcat = unions + #if __GLASGOW_HASKELL__ {-------------------------------------------------------------------- @@ -203,6 +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 f = gcast2 f #endif @@ -224,17 +253,36 @@ size t Bin sz k x l r -> sz --- | /O(log n)/. Lookup the value at a key in the map. -lookup :: Ord k => k -> Map k a -> Maybe a -lookup k t +-- | /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 + Nothing -> fail "Data.Map.lookup: Key not found" +lookup' :: Ord k => k -> Map k a -> Maybe a +lookup' k t = case t of Tip -> Nothing Bin sz kx x l r -> case compare k kx of - LT -> lookup k l - GT -> lookup k r + LT -> lookup' k l + GT -> lookup' k r EQ -> Just x +lookupAssoc :: Ord k => k -> Map k a -> Maybe (k,a) +lookupAssoc k t + = case t of + Tip -> Nothing + Bin sz kx x l r + -> case compare k kx of + LT -> lookupAssoc k l + GT -> lookupAssoc k r + EQ -> Just (kx,x) + -- | /O(log n)/. Is the key a member of the map? member :: Ord k => k -> Map k a -> Bool member k m @@ -242,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 @@ -275,9 +327,11 @@ singleton k x {-------------------------------------------------------------------- Insertion - [insert] is the inlined version of [insertWith (\k x y -> x)] --------------------------------------------------------------------} -- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value, i.e. 'insert' is equivalent to +-- @'insertWith' 'const'@. insert :: Ord k => k -> a -> Map k a -> Map k a insert kx x t = case t of @@ -289,11 +343,26 @@ insert kx x t EQ -> Bin sz kx x l r -- | /O(log n)/. Insert with a combining function. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key, f new_value old_value)@. 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 +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key,f key new_value old_value)@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a insertWithKey f kx x t = case t of @@ -302,7 +371,19 @@ insertWithKey f kx x t -> 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 -> Bin sy ky (f ky x y) l 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@) @@ -315,7 +396,7 @@ insertLookupWithKey f kx x t -> case compare kx ky of LT -> let (found,l') = insertLookupWithKey f kx x l in (found,balance ky y l' r) GT -> let (found,r') = insertLookupWithKey f kx x r in (found,balance ky y l r') - EQ -> (Just y, Bin sy ky (f ky x y) l r) + EQ -> (Just y, Bin sy kx (f kx x y) l r) {-------------------------------------------------------------------- Deletion @@ -381,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 --------------------------------------------------------------------} @@ -395,9 +493,10 @@ findIndex k t -- | /O(log n)/. Lookup the /index/ of a key. The index is a number from -- /0/ up to, but not including, the 'size' of the map. -lookupIndex :: Ord k => k -> Map k a -> Maybe Int -lookupIndex k t - = lookup 0 t +lookupIndex :: (Monad m,Ord k) => k -> Map k a -> m Int +lookupIndex k t = case lookup 0 t of + Nothing -> fail "Data.Map.lookupIndex: Key not found." + Just x -> return x where lookup idx Tip = Nothing lookup idx (Bin _ kx x l r) @@ -446,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 @@ -497,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. @@ -518,13 +644,11 @@ unionsWith f ts -- It prefers @t1@ when duplicate keys are encountered, -- i.e. (@'union' == 'unionWith' 'const'@). -- The implementation uses the efficient /hedge-union/ algorithm. --- Hedge-union is more efficient on (bigset `union` smallset)? +-- Hedge-union is more efficient on (bigset `union` smallset) union :: Ord k => Map k a -> Map k a -> Map k a union Tip t2 = t2 union t1 Tip = t1 -union t1 t2 - | size t1 >= size t2 = hedgeUnionL (const LT) (const GT) t1 t2 - | otherwise = hedgeUnionR (const LT) (const GT) t2 t1 +union t1 t2 = hedgeUnionL (const LT) (const GT) t1 t2 -- left-biased hedge union hedgeUnionL cmplo cmphi t1 Tip @@ -551,7 +675,7 @@ hedgeUnionR cmplo cmphi (Bin _ kx x l r) t2 (found,gt) = trimLookupLo kx cmphi t2 newx = case found of Nothing -> x - Just y -> y + Just (_,y) -> y {-------------------------------------------------------------------- Union with a combining function @@ -567,11 +691,7 @@ unionWith f m1 m2 unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a unionWithKey f Tip t2 = t2 unionWithKey f t1 Tip = t1 -unionWithKey f t1 t2 - | size t1 >= size t2 = hedgeUnionWithKey f (const LT) (const GT) t1 t2 - | otherwise = hedgeUnionWithKey flipf (const LT) (const GT) t2 t1 - where - flipf k x y = f k y x +unionWithKey f t1 t2 = hedgeUnionWithKey f (const LT) (const GT) t1 t2 hedgeUnionWithKey f cmplo cmphi t1 Tip = t1 @@ -586,7 +706,7 @@ hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2 (found,gt) = trimLookupLo kx cmphi t2 newx = case found of Nothing -> x - Just y -> f kx x y + Just (_,y) -> f kx x y {-------------------------------------------------------------------- Difference @@ -631,9 +751,10 @@ hedgeDiffWithKey f cmplo cmphi (Bin _ kx x l r) Tip hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r) = case found of Nothing -> merge tl tr - Just y -> case f kx y x of - Nothing -> merge tl tr - Just z -> join kx z tl tr + Just (ky,y) -> + case f ky y x of + Nothing -> merge tl tr + Just z -> join ky z tl tr where cmpkx k = compare kx k lt = trim cmplo cmpkx t @@ -659,25 +780,40 @@ intersectionWith f m1 m2 -- | /O(n+m)/. Intersection with a combining function. -- Intersection is more efficient on (bigset `intersection` smallset) +--intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c +--intersectionWithKey f Tip t = Tip +--intersectionWithKey f t Tip = Tip +--intersectionWithKey f t1 t2 = intersectWithKey f t1 t2 +-- +--intersectWithKey f Tip t = Tip +--intersectWithKey f t Tip = Tip +--intersectWithKey f t (Bin _ kx x l r) +-- = case found of +-- Nothing -> merge tl tr +-- Just y -> join kx (f kx y x) tl tr +-- where +-- (lt,found,gt) = splitLookup kx t +-- tl = intersectWithKey f lt l +-- tr = intersectWithKey f gt r + + intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c intersectionWithKey f Tip t = Tip intersectionWithKey f t Tip = Tip -intersectionWithKey f t1 t2 - | size t1 >= size t2 = intersectWithKey f t1 t2 - | otherwise = intersectWithKey flipf t2 t1 - where - flipf k x y = f k y x - -intersectWithKey f Tip t = Tip -intersectWithKey f t Tip = Tip -intersectWithKey f t (Bin _ kx x l r) - = case found of +intersectionWithKey f t1@(Bin s1 k1 x1 l1 r1) t2@(Bin s2 k2 x2 l2 r2) = + if s1 >= s2 then + let (lt,found,gt) = splitLookupWithKey k2 t1 + tl = intersectionWithKey f lt l2 + tr = intersectionWithKey f gt r2 + in case found of + Just (k,x) -> join k (f k x x2) tl tr + Nothing -> merge tl tr + else let (lt,found,gt) = splitLookup k1 t2 + tl = intersectionWithKey f l1 lt + tr = intersectionWithKey f r1 gt + in case found of + Just x -> join k1 (f k1 x1 x) tl tr Nothing -> merge tl tr - Just y -> join kx (f kx y x) tl tr - where - (lt,found,gt) = splitLookup kx t - tl = intersectWithKey f lt l - tr = intersectWithKey f gt r @@ -780,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 @@ -1058,15 +1221,15 @@ trim cmplo cmphi t@(Bin sx kx x l r) le -> trim cmplo cmphi l ge -> trim cmplo cmphi r -trimLookupLo :: Ord k => k -> (k -> Ordering) -> Map k a -> (Maybe a, Map k a) +trimLookupLo :: Ord k => k -> (k -> Ordering) -> Map k a -> (Maybe (k,a), Map k a) trimLookupLo lo cmphi Tip = (Nothing,Tip) trimLookupLo lo cmphi t@(Bin sx kx x l r) = case compare lo kx of LT -> case cmphi kx of - GT -> (lookup lo t, t) + GT -> (lookupAssoc lo t, t) le -> trimLookupLo lo cmphi l GT -> trimLookupLo lo cmphi r - EQ -> (Just x,trim (compare lo) cmphi r) + EQ -> (Just (kx,x),trim (compare lo) cmphi r) {-------------------------------------------------------------------- @@ -1112,6 +1275,22 @@ splitLookup k (Bin sx kx x l r) GT -> let (lt,z,gt) = splitLookup k r in (join kx x l lt,z,gt) EQ -> (l,Just x,r) +-- | /O(log n)/. +splitLookupWithKey :: Ord k => k -> Map k a -> (Map k a,Maybe (k,a),Map k a) +splitLookupWithKey k Tip = (Tip,Nothing,Tip) +splitLookupWithKey k (Bin sx kx x l r) + = case compare k kx of + LT -> let (lt,z,gt) = splitLookupWithKey k l in (lt,z,join kx x gt r) + GT -> let (lt,z,gt) = splitLookupWithKey k r in (join kx x l lt,z,gt) + EQ -> (l,Just (kx, x),r) + +-- | /O(log n)/. Performs a 'split' but also returns whether the pivot +-- element was found in the original set. +splitMember :: Ord k => k -> Map k a -> (Map k a,Bool,Map k a) +splitMember x t = let (l,m,r) = splitLookup x t in + (l,maybe False (const True) m,r) + + {-------------------------------------------------------------------- Utility functions that maintain the balance properties of the tree. All constructors assume that all values in [l] < [k] and all values @@ -1227,7 +1406,7 @@ deleteFindMax t - A lower [delta] leads to a more 'perfectly' balanced tree. - A higher [delta] performs less rebalancing. - - Balancing is automaic for random data and a balancing + - Balancing is automatic for random data and a balancing scheme is only necessary to avoid pathological worst cases. Almost any choice will do, and in practice, a rather large [delta] may perform better than smaller one. @@ -1290,28 +1469,55 @@ instance (Eq k,Eq a) => Eq (Map k a) where --------------------------------------------------------------------} instance (Ord k, Ord v) => Ord (Map k v) where - compare m1 m2 = compare (toList m1) (toList m2) + compare m1 m2 = compare (toAscList m1) (toAscList m2) {-------------------------------------------------------------------- - Monoid + Functor --------------------------------------------------------------------} +instance Functor (Map k) where + fmap f m = map f m -instance (Ord k) => Monoid (Map k v) where - mempty = empty - mappend = union - mconcat = unions +instance Traversable (Map k) where + traverse f Tip = pure Tip + traverse f (Bin s k v l r) + = flip (Bin s k) <$> traverse f l <*> f v <*> traverse f r + +instance Foldable (Map k) where + foldMap _f Tip = mempty + foldMap f (Bin _s _k v l r) + = foldMap f l `mappend` f v `mappend` foldMap f r {-------------------------------------------------------------------- - Functor + Read --------------------------------------------------------------------} -instance Functor (Map k) where - fmap f m = map f m +instance (Ord k, Read k, Read e) => Read (Map k e) where +#ifdef __GLASGOW_HASKELL__ + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + return (fromList xs,t) +#endif + +-- parses a pair of things with the syntax a:=b +readPair :: (Read a, Read b) => ReadS (a,b) +readPair s = do (a, ct1) <- reads s + (":=", ct2) <- lex ct1 + (b, ct3) <- reads ct2 + return ((a,b), ct3) {-------------------------------------------------------------------- Show --------------------------------------------------------------------} instance (Show k, Show a) => Show (Map k a) where - showsPrec d m = showMap (toAscList m) + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) showMap :: (Show k,Show a) => [(k,a)] -> ShowS showMap [] @@ -1320,9 +1526,9 @@ showMap (x:xs) = showChar '{' . showElem x . showTail xs where showTail [] = showChar '}' - showTail (x:xs) = showChar ',' . showElem x . showTail xs + showTail (x:xs) = showString ", " . showElem x . showTail xs - showElem (k,x) = shows k . showString ":=" . shows x + showElem (k,x) = shows k . showString " := " . shows x -- | /O(n)/. Show the tree that implements the map. The tree is shown