|
sourcecode:
|
module Data.Map (
-- Although this type is abstract, we export the constructors
-- to enable the definition of specific `Map` instances,
-- like a `RW.Base.ReadWrite` instance.
Map(..),
empty,
singleton,
fromList,
insert,
insertWith,
insertList,
insertListWith,
delete,
deleteAll,
adjust,
splitLookup,
union,
unionWith,
difference,
intersection,
intersectionWith,
foldrWithKey,
mapWithKey,
filterWithKey,
size,
null,
member,
lookup,
findWithDefault,
toList,
keys,
elems,
sortWithMap,
lookupMin,
lookupMax,
toPreOrderList
) where
import Data.Maybe
import Prelude hiding (empty)
-----------------------------------------------
-- BUILDING finite maps
-----------------------------------------------
--- The empty map.
--- @result an empty map
empty :: Map _ _
empty = Tip
--- Construct a map with only a single element.
--- @param k key of
--- @param a the single element to form
--- @result a finite map with only a single element
singleton :: k -> a -> Map k a
singleton k a = Bin k a 1 Tip Tip
--- Builts a map from given list of tuples (key,element).
--- For multiple occurences of key, the last corresponding
--- element of the list is taken.
fromList :: Ord k => [(k,a)] -> Map k a
fromList xs = insertList xs empty
-----------------------------------------------
-- ADDING AND DELETING
-----------------------------------------------
--- Throws away any previous binding and stores the new one given.
insert :: Ord k => k -> a -> Map k a -> Map k a
insert k a m = insertWith (\ _ new -> new) k a m
--- Instead of throwing away the old binding,
--- insertWith combines the new element with the old one.
--- @param combiner a function combining to elements
--- @param k the key of the elements to be combined
--- @param a the new element
--- @param m a map
--- @result a modified map
insertWith :: Ord k => (a -> a -> a)
-> k -> a -> Map k a -> Map k a
insertWith _ k a Tip = singleton k a
insertWith combiner new_k new_a (Bin k a sizeM m_l m_r)
= if new_k < k
then mkBalBranch k a (insertWith combiner new_k new_a m_l) m_r
else
if new_k == k
then Bin new_k (combiner a new_a) sizeM m_l m_r
else mkBalBranch k a m_l (insertWith combiner new_k new_a m_r)
--- Throws away any previous bindings and stores the new ones given.
--- The items are added starting with the first one in the list
insertList :: Ord k => [(k, a)] -> Map k a -> Map k a
insertList k_a_pairs m = insertListWith (\ _ new -> new) k_a_pairs m
--- Combine with a list of tuples (key, element), cf. insertWith
insertListWith :: Ord k => (a -> a -> a)
-> [(k, a)] -> Map k a -> Map k a
insertListWith combiner k_a_pairs m
= foldl add m k_a_pairs -- foldl adds from the left
where
add m' (k, a) = insertWith combiner k a m'
--- Deletes key from map.
--- Deletion doesn't complain if you try to delete something
--- which isn't there
delete :: Ord k => k -> Map k a -> Map k a
delete _ Tip = Tip
delete del_k (Bin k a _ m_l m_r)
= if del_k < k
then mkBalBranch k a (delete del_k m_l) m_r
else
if del_k == k
then glueBal m_l m_r
else mkBalBranch k a m_l (delete del_k m_r)
--- Deletes a list of keys from map.
--- Deletion doesn't complain if you try to delete something
--- which isn't there
deleteAll :: Ord k => [k] -> Map k a -> Map k a
deleteAll ks m = foldl (flip delete) m ks
--- Applies a function to element bound to given key.
adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
adjust _ _ Tip = Tip
adjust f i (Bin k x h l r)
| i == k = Bin k (f x) h l r
| i < k = Bin k x h (adjust f i l) r
| otherwise = Bin k x h l (adjust f i r)
--- Combines delFrom and lookup.
splitLookup :: Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a)
splitLookup _ Tip = (Tip, Nothing, Tip)
splitLookup i (Bin k a _ m_l m_r)
| i == k = (m_l, Just a, m_r)
| i < k = let (m_l', v, m_r') = splitLookup i m_l
in (m_l', v, glueBal m_r' m_r)
| otherwise = let (m_l', v, m_r') = splitLookup i m_r
in (glueBal m_l m_l', v, m_r')
-------------------------------------------------
-- COMBINING finite maps
-------------------------------------------------
--- Efficiently add key/element mappings of two maps into a single one.
--- CHANGED: Bindings in left argument shadow those in the right
union :: Ord k => Map k a -> Map k a -> Map k a
union Tip m2 = m2
union m1@(Bin _ _ _ _ _) Tip = m1
union (Bin split_key1 a1 _ left1 right1) m2@(Bin _ _ _ _ _)
= mkVBalBranch split_key1 a1 (union left1 lts) (union right1 gts)
where
lts = splitLT m2 split_key1
gts = splitGT m2 split_key1
--- Efficiently combine key/element mappings of two maps into a single one,
--- cf. insertWith
unionWith :: Ord k => (a -> a -> a)
-> Map k a -> Map k a -> Map k a
unionWith _ Tip m2 = m2
unionWith _ m1@(Bin _ _ _ _ _) Tip = m1
unionWith combiner (Bin split_key1 a1 _ left1 right1) m2@(Bin _ _ _ _ _)
= mkVBalBranch split_key1 new_a
(unionWith combiner left1 lts)
(unionWith combiner right1 gts)
where
lts = splitLT m2 split_key1
gts = splitGT m2 split_key1
new_a = case lookup split_key1 m2 of
Nothing -> a1
Just a2 -> combiner a2 a1
--- (difference a1 a2) deletes from a1 any bindings which are bound in a2
difference :: Ord k => Map k a -> Map k b -> Map k a
difference Tip _ = Tip
difference m1@(Bin _ _ _ _ _) Tip = m1
difference m1@(Bin _ _ _ _ _) (Bin split_key2 _ _ left2 right2)
= glueVBal (difference lts left2) (difference gts right2)
-- The two can be way different, so we need glueVBal
where
lts = splitLT m1 split_key2 -- NB gt and lt, so the equal ones
gts = splitGT m1 split_key2 -- are not in either.
--- Filters only those keys that are bound in both of the given maps.
--- CHANGED: The elements will be taken from the first map.
intersection :: Ord k => Map k a -> Map k a -> Map k a
intersection m1 m2 = intersectionWith (\ left _ -> left) m1 m2
--- Filters only those keys that are bound in both of the given maps
--- and combines the elements as in insertWith.
intersectionWith :: Ord k => (a -> b -> c)
-> Map k a -> Map k b -> Map k c
intersectionWith _ _ Tip = Tip
intersectionWith _ Tip (Bin _ _ _ _ _) = Tip
intersectionWith combiner m1@(Bin _ _ _ _ _) (Bin split_key2 a2 _ left2 right2)
| isJust maybe_a1 -- split_a *is* in intersection
= mkVBalBranch split_key2 (combiner a1' a2)
(intersectionWith combiner lts left2)
(intersectionWith combiner gts right2)
| otherwise -- split_a is *not* in intersection
= glueVBal (intersectionWith combiner lts left2)
(intersectionWith combiner gts right2)
where
lts = splitLT m1 split_key2 -- NB gt and lt, so the equal ones
gts = splitGT m1 split_key2 -- are not in either.
maybe_a1 = lookup split_key2 m1
Just a1' = maybe_a1
-------------------------------------------------------------
-- MAPPING, FOLDING, FILTERING on maps
-------------------------------------------------------------
--- Folds map by given function.
foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey _ z Tip = z
foldrWithKey k z (Bin key elt _ fm_l fm_r)
= foldrWithKey k (k key elt (foldrWithKey k z fm_r)) fm_l
--- Applies a given function on every element in the map.
mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
mapWithKey _ Tip = Tip
mapWithKey f (Bin key a size' m_l m_r)
= Bin key (f key a) size' (mapWithKey f m_l) (mapWithKey f m_r)
--- Yields a new map with only those key/element pairs matching the
--- given predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey _ Tip = Tip
filterWithKey p (Bin key a _ m_l m_r)
| p key a -- Keep the item
= mkVBalBranch key a (filterWithKey p m_l) (filterWithKey p m_r)
| otherwise -- Drop the item
= glueVBal (filterWithKey p m_l) (filterWithKey p m_r)
-----------------------------------------------------
-- INTERROGATING maps
-----------------------------------------------------
--- How many elements does given map contain?
size :: Map _ _ -> Int
size Tip = 0
size (Bin _ _ size' _ _) = size'
--- Is the given finite map empty?
null :: Map _ _ -> Bool
null m = size m == 0
--- Does given map contain given key?
member :: Ord k => k -> Map k _ -> Bool
member k m = isJust (lookup k m)
--- Retrieves element bound to given key
lookup :: Ord k => k -> Map k a -> Maybe a
lookup _ Tip = Nothing
lookup key_to_find (Bin k a _ m_l m_r)
= if key_to_find < k
then lookup key_to_find m_l
else if key_to_find == k
then Just a
else lookup key_to_find m_r
--- Retrieves element bound to given key.
--- If the element is not contained in map, return
--- default value.
findWithDefault :: Ord k => a -> k -> Map k a -> a
findWithDefault deflt k m
= case lookup k m of
Nothing -> deflt
Just a -> a
--- Retrieves the smallest key/element pair in the finite map
--- according to the basic key ordering.
lookupMin :: Map k a -> Maybe (k, a)
lookupMin Tip = Nothing
lookupMin (Bin k x _ l _) | isBranch l = lookupMin l
| otherwise = Just (k, x)
--- Retrieves the greatest key/element pair in the finite map
--- according to the basic key ordering.
lookupMax :: Map k a -> Maybe (k, a)
lookupMax Tip = Nothing
lookupMax (Bin k x _ _ r) | isBranch r = lookupMax r
| otherwise = Just (k, x)
----------------------------------------------------
-- LISTIFYING: transform finite maps to lists
----------------------------------------------------
--- Builds a list of key/element pairs. The list is ordered
--- by the "Ord" context on keys.
toList :: Map k a -> [(k, a)]
toList m = foldrWithKey (\ k a rest -> (k, a) : rest) [] m
--- Retrieves a list of keys contained in map.
--- The list is ordered
--- by the "Ord" context on keys.
keys :: Map k _ -> [k]
keys m = foldrWithKey (\ k _ rest -> k : rest) [] m
--- Retrieves a list of elements contained in map.
--- The list is ordered
--- by the "Ord" context on keys.
elems :: Map _ a -> [a]
elems m = foldrWithKey (\ _ a rest -> a : rest) [] m
--- Retrieves list of key/element pairs in preorder of the internal tree.
--- Useful for lists that will be retransformed into a tree or to match
--- any elements regardless of basic order.
toPreOrderList :: Map k a -> [(k, a)]
toPreOrderList m = pre m []
where
pre Tip xs = xs
pre (Bin k x _ l r) xs = (k, x) : pre l (pre r xs)
--- Sorts a given list by inserting and retrieving from map.
--- Duplicates are deleted.
sortWithMap :: Ord k => [k] -> [k]
sortWithMap l = keys (fromList (zip l (repeat ())))
-----------------------------------------------------
-- internal Implementation
-----------------------------------------------------
data Map k a
= Tip
| Bin k a -- Key and element stored here
Int{-STRICT-} -- Size >= 1
(Map k a) -- Children
(Map k a)
deriving (Show, Read)
instance (Eq k, Eq a) => Eq (Map k a) where
m_1 == m_2 =
(size m_1 == size m_2) && -- quick test
(toList m_1 == toList m_2)
isBranch :: Map _ _ -> Bool
isBranch (Bin _ _ _ _ _) = True
isBranch Tip = False
-------------------------------------------------------------------------
-- -
-- The implementation of balancing -
-- -
-------------------------------------------------------------------------
-------------------------------------------------------------------------
-- -
-- Basic construction of a FiniteMap -
-- -
-------------------------------------------------------------------------
sIZE_RATIO :: Int
sIZE_RATIO = 5
mkBranch :: Int
-> key -> elt
-> Map key elt -> Map key elt
-> Map key elt
mkBranch _{-which-} key elt fm_l fm_r =
let result = Bin key elt (unbox (1 + left_size + right_size)) fm_l fm_r
in
result
-- if size result <= 8 then
-- result
-- else
-- pprTrace ("mkBranch:"++(show which)) (ppr result) (
-- result
-- )
where
{-left_ok = case fm_l of
Tip -> True
Bin _ _ _ _ _ -> cmpWithBiggest_left_key key
cmpWithBiggest_left_key key' = (fst (findMax fm_l)) < key'
right_ok = case fm_r of
Tip -> True
Bin _ _ _ _ _ -> cmpWithSmallest_right_key key
cmpWithSmallest_right_key key' = key' < (fst (findMin fm_r))
balance_ok = True -- sigh-}
left_size = size fm_l
right_size = size fm_r
unbox :: Int -> Int
unbox x = x
-------------------------------------------------------------------------
-- -
-- Balanced construction of a FiniteMap -
-- -
-------------------------------------------------------------------------
mkBalBranch :: key -> elt
-> Map key elt -> Map key elt
-> Map key elt
mkBalBranch key elt fm_L fm_R
| size_l + size_r < 2
= mkBranch 1{-which-} key elt fm_L fm_R
| size_r > sIZE_RATIO * size_l -- Right tree too big
= case fm_R of
Bin _ _ _ fm_rl fm_rr ->
if size fm_rl < 2 * size fm_rr
then single_L fm_L fm_R
else double_L fm_L fm_R
-- Other case impossible
Tip -> error "Data.Map.mkBalBranch"
| size_l > sIZE_RATIO * size_r -- Left tree too big
= case fm_L of
Bin _ _ _ fm_ll fm_lr ->
if size fm_lr < 2 * size fm_ll
then single_R fm_L fm_R
else double_R fm_L fm_R
-- Other case impossible
Tip -> error "Data.Map.mkBalBranch"
| otherwise -- No imbalance
= mkBranch 2{-which-} key elt fm_L fm_R
where
size_l = size fm_L
size_r = size fm_R
single_L fm_l (Bin key_r elt_r _ fm_rl fm_rr)
= mkBranch 3{-which-} key_r elt_r (mkBranch 4{-which-} key elt fm_l fm_rl) fm_rr
single_L _ Tip = error "Data.Map.single_L"
double_L fm_l (Bin key_r elt_r _ (Bin key_rl elt_rl _ fm_rll fm_rlr) fm_rr)
= mkBranch 5{-which-} key_rl elt_rl (mkBranch 6{-which-} key elt fm_l fm_rll)
(mkBranch 7{-which-} key_r elt_r fm_rlr fm_rr)
double_L _ Tip = error "Data.Map.double_L"
double_L _ (Bin _ _ _ Tip _) = error "Data.Map.double_L"
single_R (Bin key_l elt_l _ fm_ll fm_lr) fm_r
= mkBranch 8{-which-} key_l elt_l fm_ll (mkBranch 9{-which-} key elt fm_lr fm_r)
single_R Tip _ = error "Data.Map.single_R"
double_R (Bin key_l elt_l _ fm_ll (Bin key_lr elt_lr _ fm_lrl fm_lrr)) fm_r
= mkBranch 10{-which-} key_lr elt_lr (mkBranch 11{-which-} key_l elt_l fm_ll fm_lrl)
(mkBranch 12{-which-} key elt fm_lrr fm_r)
double_R Tip _ = error "Data.Map.double_R"
double_R (Bin _ _ _ _ Tip) _ = error "Data.Map.double_R"
mkVBalBranch :: Ord key => key -> elt -> Map key elt
-> Map key elt -> Map key elt
-- Assert: in any call to (mkVBalBranch_C comb key elt l r),
-- (a) all keys in l are < all keys in r
-- (b) all keys in l are < key
-- (c) all keys in r are > key
mkVBalBranch key elt Tip fm_r = insert key elt fm_r
mkVBalBranch key elt (Bin key_l elt_l s_l fm_ll fm_lr) Tip =
insert key elt (Bin key_l elt_l s_l fm_ll fm_lr)
mkVBalBranch key elt (Bin key_l elt_l s_l fm_ll fm_lr)
(Bin key_r elt_r s_r fm_rl fm_rr)
| sIZE_RATIO * size_l < size_r
= mkBalBranch key_r elt_r (mkVBalBranch key elt fm_l fm_rl) fm_rr
| sIZE_RATIO * size_r < size_l
= mkBalBranch key_l elt_l fm_ll (mkVBalBranch key elt fm_lr fm_r)
| otherwise
= mkBranch 13{-which-} key elt fm_l fm_r
where
fm_l = Bin key_l elt_l s_l fm_ll fm_lr
fm_r = Bin key_r elt_r s_r fm_rl fm_rr
size_l = size fm_l
size_r = size fm_r
-------------------------------------------------------------------------
-- -
-- Gluing two trees together -
-- -
-------------------------------------------------------------------------
glueBal :: Map key elt -> Map key elt -> Map key elt
glueBal fm1 fm2 =
if null fm1
then fm2
else if null fm2
then fm1
else
-- The case analysis here (absent in Adams' program) is really to deal
-- with the case where fm2 is a singleton. Then deleting the minimum means
-- we pass an empty tree to mkBalBranch, which breaks its invariant.
let (mid_key1, mid_elt1) = findMax fm1
(mid_key2, mid_elt2) = findMin fm2
in
if size fm2 > size fm1
then mkBalBranch mid_key2 mid_elt2 fm1 (deleteMin fm2)
else mkBalBranch mid_key1 mid_elt1 (deleteMax fm1) fm2
glueVBal :: Map key elt -> Map key elt -> Map key elt
glueVBal fm_l fm_r =
if null fm_l
then fm_r
else if null fm_r
then fm_l
else
let Bin key_l elt_l _ fm_ll fm_lr = fm_l
Bin key_r elt_r _ fm_rl fm_rr = fm_r
--(mid_key_l,mid_elt_l) = findMax fm_l
--(mid_key_r,mid_elt_r) = findMin fm_r
size_l = size fm_l
size_r = size fm_r
in
if sIZE_RATIO * size_l < size_r
then
mkBalBranch key_r elt_r (glueVBal fm_l fm_rl) fm_rr
else if sIZE_RATIO * size_r < size_l
then
mkBalBranch key_l elt_l fm_ll (glueVBal fm_lr fm_r)
-- We now need the same two cases as in glueBal above.
else glueBal fm_l fm_r
-------------------------------------------------------------------------
-- -
-- Local utilities -
-- -
-------------------------------------------------------------------------
splitLT, splitGT :: Ord key => Map key elt -> key -> Map key elt
-- splitLT fm split_key = fm restricted to keys < split_key
-- splitGT fm split_key = fm restricted to keys > split_key
splitLT Tip _ = Tip
splitLT (Bin key elt _ fm_l fm_r) split_key
= if split_key < key
then splitLT fm_l split_key
else if split_key == key
then fm_l
else mkVBalBranch key elt fm_l (splitLT fm_r split_key)
splitGT Tip _ = Tip
splitGT (Bin key elt _ fm_l fm_r) split_key
= if split_key < key
then mkVBalBranch key elt (splitGT fm_l split_key) fm_r
else if split_key == key
then fm_r
else splitGT fm_r split_key
findMin :: Map key elt -> (key, elt)
findMin Tip = error "Data.Map.findMin: empty map"
findMin (Bin key elt _ Tip _) = (key, elt)
findMin (Bin _ _ _ (Bin key_l elt_l s_l fm_ll fm_lr)_) =
findMin (Bin key_l elt_l s_l fm_ll fm_lr)
deleteMin :: Map key elt -> Map key elt
deleteMin Tip = error "Data.Map.deleteMin: empty map"
deleteMin (Bin _ _ _ Tip fm_r) = fm_r
deleteMin (Bin key elt _ (Bin key_l elt_l s_l fm_ll fm_lr) fm_r) =
mkBalBranch key elt (deleteMin (Bin key_l elt_l s_l fm_ll fm_lr))
fm_r
findMax :: Map key elt -> (key,elt)
findMax Tip = error "Data.Map.findMax: empty map"
findMax (Bin key elt _ _ Tip) = (key, elt)
findMax (Bin _ _ _ _ (Bin key_r elt_r s_r fm_rl fm_rr)) =
findMax (Bin key_r elt_r s_r fm_rl fm_rr)
deleteMax :: Map key elt -> Map key elt
deleteMax Tip = error "FiniteMap.deleteMax: empty map"
deleteMax (Bin _ _ _ fm_l Tip) = fm_l
deleteMax (Bin key elt _ fm_l (Bin key_r elt_r s_r fm_rl fm_rr)) =
mkBalBranch key elt fm_l
(deleteMax (Bin key_r elt_r s_r fm_rl fm_rr))
------------------------------------------------------------------------------
|