----------------------------------------------------------------------------- --- A finite map is an efficient purely functional data structure --- to store a mapping from keys to values. --- --- This version was ported from a corresponding Haskell library --- --- @author Frank Huch, Bernd Brassel --- @version March 2013 --- @category algorithm ----------------------------------------------------------------------------- module Data.Map ( Map, -- abstract type 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 Maybe ----------------------------------------------- -- 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))