[project @ 2005-10-25 09:29:16 by ross]
[ghc.git] / libraries / base / Data / Map.hs
1 -----------------------------------------------------------------------------
2 -- |
3 -- Module : Data.Map
4 -- Copyright : (c) Daan Leijen 2002
5 -- License : BSD-style
6 -- Maintainer : libraries@haskell.org
7 -- Stability : provisional
8 -- Portability : portable
9 --
10 -- An efficient implementation of maps from keys to values (dictionaries).
11 --
12 -- This module is intended to be imported @qualified@, to avoid name
13 -- clashes with Prelude functions. eg.
14 --
15 -- > import Data.Map as Map
16 --
17 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
18 -- trees of /bounded balance/) as described by:
19 --
20 -- * Stephen Adams, \"/Efficient sets: a balancing act/\",
21 -- Journal of Functional Programming 3(4):553-562, October 1993,
22 -- <http://www.swiss.ai.mit.edu/~adams/BB>.
23 --
24 -- * J. Nievergelt and E.M. Reingold,
25 -- \"/Binary search trees of bounded balance/\",
26 -- SIAM journal of computing 2(1), March 1973.
27 -----------------------------------------------------------------------------
28
29 module Data.Map (
30 -- * Map type
31 Map -- instance Eq,Show,Read
32
33 -- * Operators
34 , (!), (\\)
35
36
37 -- * Query
38 , null
39 , size
40 , member
41 , lookup
42 , findWithDefault
43
44 -- * Construction
45 , empty
46 , singleton
47
48 -- ** Insertion
49 , insert
50 , insertWith, insertWithKey, insertLookupWithKey
51
52 -- ** Delete\/Update
53 , delete
54 , adjust
55 , adjustWithKey
56 , update
57 , updateWithKey
58 , updateLookupWithKey
59
60 -- * Combine
61
62 -- ** Union
63 , union
64 , unionWith
65 , unionWithKey
66 , unions
67 , unionsWith
68
69 -- ** Difference
70 , difference
71 , differenceWith
72 , differenceWithKey
73
74 -- ** Intersection
75 , intersection
76 , intersectionWith
77 , intersectionWithKey
78
79 -- * Traversal
80 -- ** Map
81 , map
82 , mapWithKey
83 , mapAccum
84 , mapAccumWithKey
85 , mapKeys
86 , mapKeysWith
87 , mapKeysMonotonic
88
89 -- ** Fold
90 , fold
91 , foldWithKey
92
93 -- * Conversion
94 , elems
95 , keys
96 , keysSet
97 , assocs
98
99 -- ** Lists
100 , toList
101 , fromList
102 , fromListWith
103 , fromListWithKey
104
105 -- ** Ordered lists
106 , toAscList
107 , fromAscList
108 , fromAscListWith
109 , fromAscListWithKey
110 , fromDistinctAscList
111
112 -- * Filter
113 , filter
114 , filterWithKey
115 , partition
116 , partitionWithKey
117
118 , split
119 , splitLookup
120
121 -- * Submap
122 , isSubmapOf, isSubmapOfBy
123 , isProperSubmapOf, isProperSubmapOfBy
124
125 -- * Indexed
126 , lookupIndex
127 , findIndex
128 , elemAt
129 , updateAt
130 , deleteAt
131
132 -- * Min\/Max
133 , findMin
134 , findMax
135 , deleteMin
136 , deleteMax
137 , deleteFindMin
138 , deleteFindMax
139 , updateMin
140 , updateMax
141 , updateMinWithKey
142 , updateMaxWithKey
143
144 -- * Debugging
145 , showTree
146 , showTreeWith
147 , valid
148 ) where
149
150 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
151 import qualified Data.Set as Set
152 import qualified Data.List as List
153 import Data.Monoid (Monoid(..))
154 import Data.Typeable
155
156 {-
157 -- for quick check
158 import qualified Prelude
159 import qualified List
160 import Debug.QuickCheck
161 import List(nub,sort)
162 -}
163
164 #if __GLASGOW_HASKELL__
165 import Text.Read
166 import Data.Generics.Basics
167 import Data.Generics.Instances
168 #endif
169
170 {--------------------------------------------------------------------
171 Operators
172 --------------------------------------------------------------------}
173 infixl 9 !,\\ --
174
175 -- | /O(log n)/. Find the value at a key.
176 -- Calls 'error' when the element can not be found.
177 (!) :: Ord k => Map k a -> k -> a
178 m ! k = find k m
179
180 -- | /O(n+m)/. See 'difference'.
181 (\\) :: Ord k => Map k a -> Map k b -> Map k a
182 m1 \\ m2 = difference m1 m2
183
184 {--------------------------------------------------------------------
185 Size balanced trees.
186 --------------------------------------------------------------------}
187 -- | A Map from keys @k@ to values @a@.
188 data Map k a = Tip
189 | Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a)
190
191 type Size = Int
192
193 instance (Ord k) => Monoid (Map k v) where
194 mempty = empty
195 mappend = union
196 mconcat = unions
197
198 #if __GLASGOW_HASKELL__
199
200 {--------------------------------------------------------------------
201 A Data instance
202 --------------------------------------------------------------------}
203
204 -- This instance preserves data abstraction at the cost of inefficiency.
205 -- We omit reflection services for the sake of data abstraction.
206
207 instance (Data k, Data a, Ord k) => Data (Map k a) where
208 gfoldl f z map = z fromList `f` (toList map)
209 toConstr _ = error "toConstr"
210 gunfold _ _ = error "gunfold"
211 dataTypeOf _ = mkNorepType "Data.Map.Map"
212
213 #endif
214
215 {--------------------------------------------------------------------
216 Query
217 --------------------------------------------------------------------}
218 -- | /O(1)/. Is the map empty?
219 null :: Map k a -> Bool
220 null t
221 = case t of
222 Tip -> True
223 Bin sz k x l r -> False
224
225 -- | /O(1)/. The number of elements in the map.
226 size :: Map k a -> Int
227 size t
228 = case t of
229 Tip -> 0
230 Bin sz k x l r -> sz
231
232
233 -- | /O(log n)/. Lookup the value at a key in the map.
234 lookup :: (Monad m,Ord k) => k -> Map k a -> m a
235 lookup k t = case lookup' k t of
236 Just x -> return x
237 Nothing -> fail "Data.Map.lookup: Key not found"
238 lookup' :: Ord k => k -> Map k a -> Maybe a
239 lookup' k t
240 = case t of
241 Tip -> Nothing
242 Bin sz kx x l r
243 -> case compare k kx of
244 LT -> lookup' k l
245 GT -> lookup' k r
246 EQ -> Just x
247
248 -- | /O(log n)/. Is the key a member of the map?
249 member :: Ord k => k -> Map k a -> Bool
250 member k m
251 = case lookup k m of
252 Nothing -> False
253 Just x -> True
254
255 -- | /O(log n)/. Find the value at a key.
256 -- Calls 'error' when the element can not be found.
257 find :: Ord k => k -> Map k a -> a
258 find k m
259 = case lookup k m of
260 Nothing -> error "Map.find: element not in the map"
261 Just x -> x
262
263 -- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
264 -- the value at key @k@ or returns @def@ when the key is not in the map.
265 findWithDefault :: Ord k => a -> k -> Map k a -> a
266 findWithDefault def k m
267 = case lookup k m of
268 Nothing -> def
269 Just x -> x
270
271
272
273 {--------------------------------------------------------------------
274 Construction
275 --------------------------------------------------------------------}
276 -- | /O(1)/. The empty map.
277 empty :: Map k a
278 empty
279 = Tip
280
281 -- | /O(1)/. A map with a single element.
282 singleton :: k -> a -> Map k a
283 singleton k x
284 = Bin 1 k x Tip Tip
285
286 {--------------------------------------------------------------------
287 Insertion
288 --------------------------------------------------------------------}
289 -- | /O(log n)/. Insert a new key and value in the map.
290 -- If the key is already present in the map, the associated value is
291 -- replaced with the supplied value, i.e. 'insert' is equivalent to
292 -- @'insertWith' 'const'@.
293 insert :: Ord k => k -> a -> Map k a -> Map k a
294 insert kx x t
295 = case t of
296 Tip -> singleton kx x
297 Bin sz ky y l r
298 -> case compare kx ky of
299 LT -> balance ky y (insert kx x l) r
300 GT -> balance ky y l (insert kx x r)
301 EQ -> Bin sz kx x l r
302
303 -- | /O(log n)/. Insert with a combining function.
304 insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
305 insertWith f k x m
306 = insertWithKey (\k x y -> f x y) k x m
307
308 -- | /O(log n)/. Insert with a combining function.
309 insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
310 insertWithKey f kx x t
311 = case t of
312 Tip -> singleton kx x
313 Bin sy ky y l r
314 -> case compare kx ky of
315 LT -> balance ky y (insertWithKey f kx x l) r
316 GT -> balance ky y l (insertWithKey f kx x r)
317 EQ -> Bin sy ky (f ky x y) l r
318
319 -- | /O(log n)/. The expression (@'insertLookupWithKey' f k x map@)
320 -- is a pair where the first element is equal to (@'lookup' k map@)
321 -- and the second element equal to (@'insertWithKey' f k x map@).
322 insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a,Map k a)
323 insertLookupWithKey f kx x t
324 = case t of
325 Tip -> (Nothing, singleton kx x)
326 Bin sy ky y l r
327 -> case compare kx ky of
328 LT -> let (found,l') = insertLookupWithKey f kx x l in (found,balance ky y l' r)
329 GT -> let (found,r') = insertLookupWithKey f kx x r in (found,balance ky y l r')
330 EQ -> (Just y, Bin sy ky (f ky x y) l r)
331
332 {--------------------------------------------------------------------
333 Deletion
334 [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
335 --------------------------------------------------------------------}
336 -- | /O(log n)/. Delete a key and its value from the map. When the key is not
337 -- a member of the map, the original map is returned.
338 delete :: Ord k => k -> Map k a -> Map k a
339 delete k t
340 = case t of
341 Tip -> Tip
342 Bin sx kx x l r
343 -> case compare k kx of
344 LT -> balance kx x (delete k l) r
345 GT -> balance kx x l (delete k r)
346 EQ -> glue l r
347
348 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
349 -- a member of the map, the original map is returned.
350 adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
351 adjust f k m
352 = adjustWithKey (\k x -> f x) k m
353
354 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
355 -- a member of the map, the original map is returned.
356 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
357 adjustWithKey f k m
358 = updateWithKey (\k x -> Just (f k x)) k m
359
360 -- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@
361 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
362 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
363 update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
364 update f k m
365 = updateWithKey (\k x -> f x) k m
366
367 -- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the
368 -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing',
369 -- the element is deleted. If it is (@'Just' y@), the key @k@ is bound
370 -- to the new value @y@.
371 updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
372 updateWithKey f k t
373 = case t of
374 Tip -> Tip
375 Bin sx kx x l r
376 -> case compare k kx of
377 LT -> balance kx x (updateWithKey f k l) r
378 GT -> balance kx x l (updateWithKey f k r)
379 EQ -> case f kx x of
380 Just x' -> Bin sx kx x' l r
381 Nothing -> glue l r
382
383 -- | /O(log n)/. Lookup and update.
384 updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)
385 updateLookupWithKey f k t
386 = case t of
387 Tip -> (Nothing,Tip)
388 Bin sx kx x l r
389 -> case compare k kx of
390 LT -> let (found,l') = updateLookupWithKey f k l in (found,balance kx x l' r)
391 GT -> let (found,r') = updateLookupWithKey f k r in (found,balance kx x l r')
392 EQ -> case f kx x of
393 Just x' -> (Just x',Bin sx kx x' l r)
394 Nothing -> (Just x,glue l r)
395
396 {--------------------------------------------------------------------
397 Indexing
398 --------------------------------------------------------------------}
399 -- | /O(log n)/. Return the /index/ of a key. The index is a number from
400 -- /0/ up to, but not including, the 'size' of the map. Calls 'error' when
401 -- the key is not a 'member' of the map.
402 findIndex :: Ord k => k -> Map k a -> Int
403 findIndex k t
404 = case lookupIndex k t of
405 Nothing -> error "Map.findIndex: element is not in the map"
406 Just idx -> idx
407
408 -- | /O(log n)/. Lookup the /index/ of a key. The index is a number from
409 -- /0/ up to, but not including, the 'size' of the map.
410 lookupIndex :: (Monad m,Ord k) => k -> Map k a -> m Int
411 lookupIndex k t = case lookup 0 t of
412 Nothing -> fail "Data.Map.lookupIndex: Key not found."
413 Just x -> return x
414 where
415 lookup idx Tip = Nothing
416 lookup idx (Bin _ kx x l r)
417 = case compare k kx of
418 LT -> lookup idx l
419 GT -> lookup (idx + size l + 1) r
420 EQ -> Just (idx + size l)
421
422 -- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an
423 -- invalid index is used.
424 elemAt :: Int -> Map k a -> (k,a)
425 elemAt i Tip = error "Map.elemAt: index out of range"
426 elemAt i (Bin _ kx x l r)
427 = case compare i sizeL of
428 LT -> elemAt i l
429 GT -> elemAt (i-sizeL-1) r
430 EQ -> (kx,x)
431 where
432 sizeL = size l
433
434 -- | /O(log n)/. Update the element at /index/. Calls 'error' when an
435 -- invalid index is used.
436 updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
437 updateAt f i Tip = error "Map.updateAt: index out of range"
438 updateAt f i (Bin sx kx x l r)
439 = case compare i sizeL of
440 LT -> updateAt f i l
441 GT -> updateAt f (i-sizeL-1) r
442 EQ -> case f kx x of
443 Just x' -> Bin sx kx x' l r
444 Nothing -> glue l r
445 where
446 sizeL = size l
447
448 -- | /O(log n)/. Delete the element at /index/.
449 -- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@).
450 deleteAt :: Int -> Map k a -> Map k a
451 deleteAt i map
452 = updateAt (\k x -> Nothing) i map
453
454
455 {--------------------------------------------------------------------
456 Minimal, Maximal
457 --------------------------------------------------------------------}
458 -- | /O(log n)/. The minimal key of the map.
459 findMin :: Map k a -> (k,a)
460 findMin (Bin _ kx x Tip r) = (kx,x)
461 findMin (Bin _ kx x l r) = findMin l
462 findMin Tip = error "Map.findMin: empty tree has no minimal element"
463
464 -- | /O(log n)/. The maximal key of the map.
465 findMax :: Map k a -> (k,a)
466 findMax (Bin _ kx x l Tip) = (kx,x)
467 findMax (Bin _ kx x l r) = findMax r
468 findMax Tip = error "Map.findMax: empty tree has no maximal element"
469
470 -- | /O(log n)/. Delete the minimal key.
471 deleteMin :: Map k a -> Map k a
472 deleteMin (Bin _ kx x Tip r) = r
473 deleteMin (Bin _ kx x l r) = balance kx x (deleteMin l) r
474 deleteMin Tip = Tip
475
476 -- | /O(log n)/. Delete the maximal key.
477 deleteMax :: Map k a -> Map k a
478 deleteMax (Bin _ kx x l Tip) = l
479 deleteMax (Bin _ kx x l r) = balance kx x l (deleteMax r)
480 deleteMax Tip = Tip
481
482 -- | /O(log n)/. Update the value at the minimal key.
483 updateMin :: (a -> Maybe a) -> Map k a -> Map k a
484 updateMin f m
485 = updateMinWithKey (\k x -> f x) m
486
487 -- | /O(log n)/. Update the value at the maximal key.
488 updateMax :: (a -> Maybe a) -> Map k a -> Map k a
489 updateMax f m
490 = updateMaxWithKey (\k x -> f x) m
491
492
493 -- | /O(log n)/. Update the value at the minimal key.
494 updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
495 updateMinWithKey f t
496 = case t of
497 Bin sx kx x Tip r -> case f kx x of
498 Nothing -> r
499 Just x' -> Bin sx kx x' Tip r
500 Bin sx kx x l r -> balance kx x (updateMinWithKey f l) r
501 Tip -> Tip
502
503 -- | /O(log n)/. Update the value at the maximal key.
504 updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
505 updateMaxWithKey f t
506 = case t of
507 Bin sx kx x l Tip -> case f kx x of
508 Nothing -> l
509 Just x' -> Bin sx kx x' l Tip
510 Bin sx kx x l r -> balance kx x l (updateMaxWithKey f r)
511 Tip -> Tip
512
513
514 {--------------------------------------------------------------------
515 Union.
516 --------------------------------------------------------------------}
517 -- | The union of a list of maps:
518 -- (@'unions' == 'Prelude.foldl' 'union' 'empty'@).
519 unions :: Ord k => [Map k a] -> Map k a
520 unions ts
521 = foldlStrict union empty ts
522
523 -- | The union of a list of maps, with a combining operation:
524 -- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
525 unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a
526 unionsWith f ts
527 = foldlStrict (unionWith f) empty ts
528
529 -- | /O(n+m)/.
530 -- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@.
531 -- It prefers @t1@ when duplicate keys are encountered,
532 -- i.e. (@'union' == 'unionWith' 'const'@).
533 -- The implementation uses the efficient /hedge-union/ algorithm.
534 -- Hedge-union is more efficient on (bigset `union` smallset)?
535 union :: Ord k => Map k a -> Map k a -> Map k a
536 union Tip t2 = t2
537 union t1 Tip = t1
538 union t1 t2
539 | size t1 >= size t2 = hedgeUnionL (const LT) (const GT) t1 t2
540 | otherwise = hedgeUnionR (const LT) (const GT) t2 t1
541
542 -- left-biased hedge union
543 hedgeUnionL cmplo cmphi t1 Tip
544 = t1
545 hedgeUnionL cmplo cmphi Tip (Bin _ kx x l r)
546 = join kx x (filterGt cmplo l) (filterLt cmphi r)
547 hedgeUnionL cmplo cmphi (Bin _ kx x l r) t2
548 = join kx x (hedgeUnionL cmplo cmpkx l (trim cmplo cmpkx t2))
549 (hedgeUnionL cmpkx cmphi r (trim cmpkx cmphi t2))
550 where
551 cmpkx k = compare kx k
552
553 -- right-biased hedge union
554 hedgeUnionR cmplo cmphi t1 Tip
555 = t1
556 hedgeUnionR cmplo cmphi Tip (Bin _ kx x l r)
557 = join kx x (filterGt cmplo l) (filterLt cmphi r)
558 hedgeUnionR cmplo cmphi (Bin _ kx x l r) t2
559 = join kx newx (hedgeUnionR cmplo cmpkx l lt)
560 (hedgeUnionR cmpkx cmphi r gt)
561 where
562 cmpkx k = compare kx k
563 lt = trim cmplo cmpkx t2
564 (found,gt) = trimLookupLo kx cmphi t2
565 newx = case found of
566 Nothing -> x
567 Just y -> y
568
569 {--------------------------------------------------------------------
570 Union with a combining function
571 --------------------------------------------------------------------}
572 -- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
573 unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
574 unionWith f m1 m2
575 = unionWithKey (\k x y -> f x y) m1 m2
576
577 -- | /O(n+m)/.
578 -- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
579 -- Hedge-union is more efficient on (bigset `union` smallset).
580 unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
581 unionWithKey f Tip t2 = t2
582 unionWithKey f t1 Tip = t1
583 unionWithKey f t1 t2
584 | size t1 >= size t2 = hedgeUnionWithKey f (const LT) (const GT) t1 t2
585 | otherwise = hedgeUnionWithKey flipf (const LT) (const GT) t2 t1
586 where
587 flipf k x y = f k y x
588
589 hedgeUnionWithKey f cmplo cmphi t1 Tip
590 = t1
591 hedgeUnionWithKey f cmplo cmphi Tip (Bin _ kx x l r)
592 = join kx x (filterGt cmplo l) (filterLt cmphi r)
593 hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2
594 = join kx newx (hedgeUnionWithKey f cmplo cmpkx l lt)
595 (hedgeUnionWithKey f cmpkx cmphi r gt)
596 where
597 cmpkx k = compare kx k
598 lt = trim cmplo cmpkx t2
599 (found,gt) = trimLookupLo kx cmphi t2
600 newx = case found of
601 Nothing -> x
602 Just y -> f kx x y
603
604 {--------------------------------------------------------------------
605 Difference
606 --------------------------------------------------------------------}
607 -- | /O(n+m)/. Difference of two maps.
608 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
609 difference :: Ord k => Map k a -> Map k b -> Map k a
610 difference Tip t2 = Tip
611 difference t1 Tip = t1
612 difference t1 t2 = hedgeDiff (const LT) (const GT) t1 t2
613
614 hedgeDiff cmplo cmphi Tip t
615 = Tip
616 hedgeDiff cmplo cmphi (Bin _ kx x l r) Tip
617 = join kx x (filterGt cmplo l) (filterLt cmphi r)
618 hedgeDiff cmplo cmphi t (Bin _ kx x l r)
619 = merge (hedgeDiff cmplo cmpkx (trim cmplo cmpkx t) l)
620 (hedgeDiff cmpkx cmphi (trim cmpkx cmphi t) r)
621 where
622 cmpkx k = compare kx k
623
624 -- | /O(n+m)/. Difference with a combining function.
625 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
626 differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
627 differenceWith f m1 m2
628 = differenceWithKey (\k x y -> f x y) m1 m2
629
630 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
631 -- encountered, the combining function is applied to the key and both values.
632 -- If it returns 'Nothing', the element is discarded (proper set difference). If
633 -- it returns (@'Just' y@), the element is updated with a new value @y@.
634 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
635 differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
636 differenceWithKey f Tip t2 = Tip
637 differenceWithKey f t1 Tip = t1
638 differenceWithKey f t1 t2 = hedgeDiffWithKey f (const LT) (const GT) t1 t2
639
640 hedgeDiffWithKey f cmplo cmphi Tip t
641 = Tip
642 hedgeDiffWithKey f cmplo cmphi (Bin _ kx x l r) Tip
643 = join kx x (filterGt cmplo l) (filterLt cmphi r)
644 hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r)
645 = case found of
646 Nothing -> merge tl tr
647 Just y -> case f kx y x of
648 Nothing -> merge tl tr
649 Just z -> join kx z tl tr
650 where
651 cmpkx k = compare kx k
652 lt = trim cmplo cmpkx t
653 (found,gt) = trimLookupLo kx cmphi t
654 tl = hedgeDiffWithKey f cmplo cmpkx lt l
655 tr = hedgeDiffWithKey f cmpkx cmphi gt r
656
657
658
659 {--------------------------------------------------------------------
660 Intersection
661 --------------------------------------------------------------------}
662 -- | /O(n+m)/. Intersection of two maps. The values in the first
663 -- map are returned, i.e. (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
664 intersection :: Ord k => Map k a -> Map k b -> Map k a
665 intersection m1 m2
666 = intersectionWithKey (\k x y -> x) m1 m2
667
668 -- | /O(n+m)/. Intersection with a combining function.
669 intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
670 intersectionWith f m1 m2
671 = intersectionWithKey (\k x y -> f x y) m1 m2
672
673 -- | /O(n+m)/. Intersection with a combining function.
674 -- Intersection is more efficient on (bigset `intersection` smallset)
675 intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
676 intersectionWithKey f Tip t = Tip
677 intersectionWithKey f t Tip = Tip
678 intersectionWithKey f t1 t2
679 | size t1 >= size t2 = intersectWithKey f t1 t2
680 | otherwise = intersectWithKey flipf t2 t1
681 where
682 flipf k x y = f k y x
683
684 intersectWithKey f Tip t = Tip
685 intersectWithKey f t Tip = Tip
686 intersectWithKey f t (Bin _ kx x l r)
687 = case found of
688 Nothing -> merge tl tr
689 Just y -> join kx (f kx y x) tl tr
690 where
691 (lt,found,gt) = splitLookup kx t
692 tl = intersectWithKey f lt l
693 tr = intersectWithKey f gt r
694
695
696
697 {--------------------------------------------------------------------
698 Submap
699 --------------------------------------------------------------------}
700 -- | /O(n+m)/.
701 -- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
702 isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
703 isSubmapOf m1 m2
704 = isSubmapOfBy (==) m1 m2
705
706 {- | /O(n+m)/.
707 The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
708 all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
709 applied to their respective values. For example, the following
710 expressions are all 'True':
711
712 > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
713 > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
714 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])
715
716 But the following are all 'False':
717
718 > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
719 > isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
720 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
721 -}
722 isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool
723 isSubmapOfBy f t1 t2
724 = (size t1 <= size t2) && (submap' f t1 t2)
725
726 submap' f Tip t = True
727 submap' f t Tip = False
728 submap' f (Bin _ kx x l r) t
729 = case found of
730 Nothing -> False
731 Just y -> f x y && submap' f l lt && submap' f r gt
732 where
733 (lt,found,gt) = splitLookup kx t
734
735 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
736 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
737 isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
738 isProperSubmapOf m1 m2
739 = isProperSubmapOfBy (==) m1 m2
740
741 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
742 The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
743 @m1@ and @m2@ are not equal,
744 all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
745 applied to their respective values. For example, the following
746 expressions are all 'True':
747
748 > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
749 > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
750
751 But the following are all 'False':
752
753 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
754 > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
755 > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
756 -}
757 isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
758 isProperSubmapOfBy f t1 t2
759 = (size t1 < size t2) && (submap' f t1 t2)
760
761 {--------------------------------------------------------------------
762 Filter and partition
763 --------------------------------------------------------------------}
764 -- | /O(n)/. Filter all values that satisfy the predicate.
765 filter :: Ord k => (a -> Bool) -> Map k a -> Map k a
766 filter p m
767 = filterWithKey (\k x -> p x) m
768
769 -- | /O(n)/. Filter all keys\/values that satisfy the predicate.
770 filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
771 filterWithKey p Tip = Tip
772 filterWithKey p (Bin _ kx x l r)
773 | p kx x = join kx x (filterWithKey p l) (filterWithKey p r)
774 | otherwise = merge (filterWithKey p l) (filterWithKey p r)
775
776
777 -- | /O(n)/. partition the map according to a predicate. The first
778 -- map contains all elements that satisfy the predicate, the second all
779 -- elements that fail the predicate. See also 'split'.
780 partition :: Ord k => (a -> Bool) -> Map k a -> (Map k a,Map k a)
781 partition p m
782 = partitionWithKey (\k x -> p x) m
783
784 -- | /O(n)/. partition the map according to a predicate. The first
785 -- map contains all elements that satisfy the predicate, the second all
786 -- elements that fail the predicate. See also 'split'.
787 partitionWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
788 partitionWithKey p Tip = (Tip,Tip)
789 partitionWithKey p (Bin _ kx x l r)
790 | p kx x = (join kx x l1 r1,merge l2 r2)
791 | otherwise = (merge l1 r1,join kx x l2 r2)
792 where
793 (l1,l2) = partitionWithKey p l
794 (r1,r2) = partitionWithKey p r
795
796
797 {--------------------------------------------------------------------
798 Mapping
799 --------------------------------------------------------------------}
800 -- | /O(n)/. Map a function over all values in the map.
801 map :: (a -> b) -> Map k a -> Map k b
802 map f m
803 = mapWithKey (\k x -> f x) m
804
805 -- | /O(n)/. Map a function over all values in the map.
806 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
807 mapWithKey f Tip = Tip
808 mapWithKey f (Bin sx kx x l r)
809 = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
810
811 -- | /O(n)/. The function 'mapAccum' threads an accumulating
812 -- argument through the map in ascending order of keys.
813 mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
814 mapAccum f a m
815 = mapAccumWithKey (\a k x -> f a x) a m
816
817 -- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating
818 -- argument through the map in ascending order of keys.
819 mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
820 mapAccumWithKey f a t
821 = mapAccumL f a t
822
823 -- | /O(n)/. The function 'mapAccumL' threads an accumulating
824 -- argument throught the map in ascending order of keys.
825 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
826 mapAccumL f a t
827 = case t of
828 Tip -> (a,Tip)
829 Bin sx kx x l r
830 -> let (a1,l') = mapAccumL f a l
831 (a2,x') = f a1 kx x
832 (a3,r') = mapAccumL f a2 r
833 in (a3,Bin sx kx x' l' r')
834
835 -- | /O(n)/. The function 'mapAccumR' threads an accumulating
836 -- argument throught the map in descending order of keys.
837 mapAccumR :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
838 mapAccumR f a t
839 = case t of
840 Tip -> (a,Tip)
841 Bin sx kx x l r
842 -> let (a1,r') = mapAccumR f a r
843 (a2,x') = f a1 kx x
844 (a3,l') = mapAccumR f a2 l
845 in (a3,Bin sx kx x' l' r')
846
847 -- | /O(n*log n)/.
848 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
849 --
850 -- The size of the result may be smaller if @f@ maps two or more distinct
851 -- keys to the same new key. In this case the value at the smallest of
852 -- these keys is retained.
853
854 mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
855 mapKeys = mapKeysWith (\x y->x)
856
857 -- | /O(n*log n)/.
858 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
859 --
860 -- The size of the result may be smaller if @f@ maps two or more distinct
861 -- keys to the same new key. In this case the associated values will be
862 -- combined using @c@.
863
864 mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
865 mapKeysWith c f = fromListWith c . List.map fFirst . toList
866 where fFirst (x,y) = (f x, y)
867
868
869 -- | /O(n)/.
870 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
871 -- is strictly monotonic.
872 -- /The precondition is not checked./
873 -- Semi-formally, we have:
874 --
875 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
876 -- > ==> mapKeysMonotonic f s == mapKeys f s
877 -- > where ls = keys s
878
879 mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a
880 mapKeysMonotonic f Tip = Tip
881 mapKeysMonotonic f (Bin sz k x l r) =
882 Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r)
883
884 {--------------------------------------------------------------------
885 Folds
886 --------------------------------------------------------------------}
887
888 -- | /O(n)/. Fold the values in the map, such that
889 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
890 -- For example,
891 --
892 -- > elems map = fold (:) [] map
893 --
894 fold :: (a -> b -> b) -> b -> Map k a -> b
895 fold f z m
896 = foldWithKey (\k x z -> f x z) z m
897
898 -- | /O(n)/. Fold the keys and values in the map, such that
899 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
900 -- For example,
901 --
902 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
903 --
904 foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
905 foldWithKey f z t
906 = foldr f z t
907
908 -- | /O(n)/. In-order fold.
909 foldi :: (k -> a -> b -> b -> b) -> b -> Map k a -> b
910 foldi f z Tip = z
911 foldi f z (Bin _ kx x l r) = f kx x (foldi f z l) (foldi f z r)
912
913 -- | /O(n)/. Post-order fold.
914 foldr :: (k -> a -> b -> b) -> b -> Map k a -> b
915 foldr f z Tip = z
916 foldr f z (Bin _ kx x l r) = foldr f (f kx x (foldr f z r)) l
917
918 -- | /O(n)/. Pre-order fold.
919 foldl :: (b -> k -> a -> b) -> b -> Map k a -> b
920 foldl f z Tip = z
921 foldl f z (Bin _ kx x l r) = foldl f (f (foldl f z l) kx x) r
922
923 {--------------------------------------------------------------------
924 List variations
925 --------------------------------------------------------------------}
926 -- | /O(n)/.
927 -- Return all elements of the map in the ascending order of their keys.
928 elems :: Map k a -> [a]
929 elems m
930 = [x | (k,x) <- assocs m]
931
932 -- | /O(n)/. Return all keys of the map in ascending order.
933 keys :: Map k a -> [k]
934 keys m
935 = [k | (k,x) <- assocs m]
936
937 -- | /O(n)/. The set of all keys of the map.
938 keysSet :: Map k a -> Set.Set k
939 keysSet m = Set.fromDistinctAscList (keys m)
940
941 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
942 assocs :: Map k a -> [(k,a)]
943 assocs m
944 = toList m
945
946 {--------------------------------------------------------------------
947 Lists
948 use [foldlStrict] to reduce demand on the control-stack
949 --------------------------------------------------------------------}
950 -- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'.
951 fromList :: Ord k => [(k,a)] -> Map k a
952 fromList xs
953 = foldlStrict ins empty xs
954 where
955 ins t (k,x) = insert k x t
956
957 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
958 fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
959 fromListWith f xs
960 = fromListWithKey (\k x y -> f x y) xs
961
962 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
963 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
964 fromListWithKey f xs
965 = foldlStrict ins empty xs
966 where
967 ins t (k,x) = insertWithKey f k x t
968
969 -- | /O(n)/. Convert to a list of key\/value pairs.
970 toList :: Map k a -> [(k,a)]
971 toList t = toAscList t
972
973 -- | /O(n)/. Convert to an ascending list.
974 toAscList :: Map k a -> [(k,a)]
975 toAscList t = foldr (\k x xs -> (k,x):xs) [] t
976
977 -- | /O(n)/.
978 toDescList :: Map k a -> [(k,a)]
979 toDescList t = foldl (\xs k x -> (k,x):xs) [] t
980
981
982 {--------------------------------------------------------------------
983 Building trees from ascending/descending lists can be done in linear time.
984
985 Note that if [xs] is ascending that:
986 fromAscList xs == fromList xs
987 fromAscListWith f xs == fromListWith f xs
988 --------------------------------------------------------------------}
989 -- | /O(n)/. Build a map from an ascending list in linear time.
990 -- /The precondition (input list is ascending) is not checked./
991 fromAscList :: Eq k => [(k,a)] -> Map k a
992 fromAscList xs
993 = fromAscListWithKey (\k x y -> x) xs
994
995 -- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
996 -- /The precondition (input list is ascending) is not checked./
997 fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
998 fromAscListWith f xs
999 = fromAscListWithKey (\k x y -> f x y) xs
1000
1001 -- | /O(n)/. Build a map from an ascending list in linear time with a
1002 -- combining function for equal keys.
1003 -- /The precondition (input list is ascending) is not checked./
1004 fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1005 fromAscListWithKey f xs
1006 = fromDistinctAscList (combineEq f xs)
1007 where
1008 -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1009 combineEq f xs
1010 = case xs of
1011 [] -> []
1012 [x] -> [x]
1013 (x:xx) -> combineEq' x xx
1014
1015 combineEq' z [] = [z]
1016 combineEq' z@(kz,zz) (x@(kx,xx):xs)
1017 | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs
1018 | otherwise = z:combineEq' x xs
1019
1020
1021 -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
1022 -- /The precondition is not checked./
1023 fromDistinctAscList :: [(k,a)] -> Map k a
1024 fromDistinctAscList xs
1025 = build const (length xs) xs
1026 where
1027 -- 1) use continutations so that we use heap space instead of stack space.
1028 -- 2) special case for n==5 to build bushier trees.
1029 build c 0 xs = c Tip xs
1030 build c 5 xs = case xs of
1031 ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx)
1032 -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx
1033 build c n xs = seq nr $ build (buildR nr c) nl xs
1034 where
1035 nl = n `div` 2
1036 nr = n - nl - 1
1037
1038 buildR n c l ((k,x):ys) = build (buildB l k x c) n ys
1039 buildB l k x c r zs = c (bin k x l r) zs
1040
1041
1042
1043 {--------------------------------------------------------------------
1044 Utility functions that return sub-ranges of the original
1045 tree. Some functions take a comparison function as argument to
1046 allow comparisons against infinite values. A function [cmplo k]
1047 should be read as [compare lo k].
1048
1049 [trim cmplo cmphi t] A tree that is either empty or where [cmplo k == LT]
1050 and [cmphi k == GT] for the key [k] of the root.
1051 [filterGt cmp t] A tree where for all keys [k]. [cmp k == LT]
1052 [filterLt cmp t] A tree where for all keys [k]. [cmp k == GT]
1053
1054 [split k t] Returns two trees [l] and [r] where all keys
1055 in [l] are <[k] and all keys in [r] are >[k].
1056 [splitLookup k t] Just like [split] but also returns whether [k]
1057 was found in the tree.
1058 --------------------------------------------------------------------}
1059
1060 {--------------------------------------------------------------------
1061 [trim lo hi t] trims away all subtrees that surely contain no
1062 values between the range [lo] to [hi]. The returned tree is either
1063 empty or the key of the root is between @lo@ and @hi@.
1064 --------------------------------------------------------------------}
1065 trim :: (k -> Ordering) -> (k -> Ordering) -> Map k a -> Map k a
1066 trim cmplo cmphi Tip = Tip
1067 trim cmplo cmphi t@(Bin sx kx x l r)
1068 = case cmplo kx of
1069 LT -> case cmphi kx of
1070 GT -> t
1071 le -> trim cmplo cmphi l
1072 ge -> trim cmplo cmphi r
1073
1074 trimLookupLo :: Ord k => k -> (k -> Ordering) -> Map k a -> (Maybe a, Map k a)
1075 trimLookupLo lo cmphi Tip = (Nothing,Tip)
1076 trimLookupLo lo cmphi t@(Bin sx kx x l r)
1077 = case compare lo kx of
1078 LT -> case cmphi kx of
1079 GT -> (lookup lo t, t)
1080 le -> trimLookupLo lo cmphi l
1081 GT -> trimLookupLo lo cmphi r
1082 EQ -> (Just x,trim (compare lo) cmphi r)
1083
1084
1085 {--------------------------------------------------------------------
1086 [filterGt k t] filter all keys >[k] from tree [t]
1087 [filterLt k t] filter all keys <[k] from tree [t]
1088 --------------------------------------------------------------------}
1089 filterGt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1090 filterGt cmp Tip = Tip
1091 filterGt cmp (Bin sx kx x l r)
1092 = case cmp kx of
1093 LT -> join kx x (filterGt cmp l) r
1094 GT -> filterGt cmp r
1095 EQ -> r
1096
1097 filterLt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1098 filterLt cmp Tip = Tip
1099 filterLt cmp (Bin sx kx x l r)
1100 = case cmp kx of
1101 LT -> filterLt cmp l
1102 GT -> join kx x l (filterLt cmp r)
1103 EQ -> l
1104
1105 {--------------------------------------------------------------------
1106 Split
1107 --------------------------------------------------------------------}
1108 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where
1109 -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1110 split :: Ord k => k -> Map k a -> (Map k a,Map k a)
1111 split k Tip = (Tip,Tip)
1112 split k (Bin sx kx x l r)
1113 = case compare k kx of
1114 LT -> let (lt,gt) = split k l in (lt,join kx x gt r)
1115 GT -> let (lt,gt) = split k r in (join kx x l lt,gt)
1116 EQ -> (l,r)
1117
1118 -- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
1119 -- like 'split' but also returns @'lookup' k map@.
1120 splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a)
1121 splitLookup k Tip = (Tip,Nothing,Tip)
1122 splitLookup k (Bin sx kx x l r)
1123 = case compare k kx of
1124 LT -> let (lt,z,gt) = splitLookup k l in (lt,z,join kx x gt r)
1125 GT -> let (lt,z,gt) = splitLookup k r in (join kx x l lt,z,gt)
1126 EQ -> (l,Just x,r)
1127
1128 {--------------------------------------------------------------------
1129 Utility functions that maintain the balance properties of the tree.
1130 All constructors assume that all values in [l] < [k] and all values
1131 in [r] > [k], and that [l] and [r] are valid trees.
1132
1133 In order of sophistication:
1134 [Bin sz k x l r] The type constructor.
1135 [bin k x l r] Maintains the correct size, assumes that both [l]
1136 and [r] are balanced with respect to each other.
1137 [balance k x l r] Restores the balance and size.
1138 Assumes that the original tree was balanced and
1139 that [l] or [r] has changed by at most one element.
1140 [join k x l r] Restores balance and size.
1141
1142 Furthermore, we can construct a new tree from two trees. Both operations
1143 assume that all values in [l] < all values in [r] and that [l] and [r]
1144 are valid:
1145 [glue l r] Glues [l] and [r] together. Assumes that [l] and
1146 [r] are already balanced with respect to each other.
1147 [merge l r] Merges two trees and restores balance.
1148
1149 Note: in contrast to Adam's paper, we use (<=) comparisons instead
1150 of (<) comparisons in [join], [merge] and [balance].
1151 Quickcheck (on [difference]) showed that this was necessary in order
1152 to maintain the invariants. It is quite unsatisfactory that I haven't
1153 been able to find out why this is actually the case! Fortunately, it
1154 doesn't hurt to be a bit more conservative.
1155 --------------------------------------------------------------------}
1156
1157 {--------------------------------------------------------------------
1158 Join
1159 --------------------------------------------------------------------}
1160 join :: Ord k => k -> a -> Map k a -> Map k a -> Map k a
1161 join kx x Tip r = insertMin kx x r
1162 join kx x l Tip = insertMax kx x l
1163 join kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
1164 | delta*sizeL <= sizeR = balance kz z (join kx x l lz) rz
1165 | delta*sizeR <= sizeL = balance ky y ly (join kx x ry r)
1166 | otherwise = bin kx x l r
1167
1168
1169 -- insertMin and insertMax don't perform potentially expensive comparisons.
1170 insertMax,insertMin :: k -> a -> Map k a -> Map k a
1171 insertMax kx x t
1172 = case t of
1173 Tip -> singleton kx x
1174 Bin sz ky y l r
1175 -> balance ky y l (insertMax kx x r)
1176
1177 insertMin kx x t
1178 = case t of
1179 Tip -> singleton kx x
1180 Bin sz ky y l r
1181 -> balance ky y (insertMin kx x l) r
1182
1183 {--------------------------------------------------------------------
1184 [merge l r]: merges two trees.
1185 --------------------------------------------------------------------}
1186 merge :: Map k a -> Map k a -> Map k a
1187 merge Tip r = r
1188 merge l Tip = l
1189 merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
1190 | delta*sizeL <= sizeR = balance ky y (merge l ly) ry
1191 | delta*sizeR <= sizeL = balance kx x lx (merge rx r)
1192 | otherwise = glue l r
1193
1194 {--------------------------------------------------------------------
1195 [glue l r]: glues two trees together.
1196 Assumes that [l] and [r] are already balanced with respect to each other.
1197 --------------------------------------------------------------------}
1198 glue :: Map k a -> Map k a -> Map k a
1199 glue Tip r = r
1200 glue l Tip = l
1201 glue l r
1202 | size l > size r = let ((km,m),l') = deleteFindMax l in balance km m l' r
1203 | otherwise = let ((km,m),r') = deleteFindMin r in balance km m l r'
1204
1205
1206 -- | /O(log n)/. Delete and find the minimal element.
1207 deleteFindMin :: Map k a -> ((k,a),Map k a)
1208 deleteFindMin t
1209 = case t of
1210 Bin _ k x Tip r -> ((k,x),r)
1211 Bin _ k x l r -> let (km,l') = deleteFindMin l in (km,balance k x l' r)
1212 Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip)
1213
1214 -- | /O(log n)/. Delete and find the maximal element.
1215 deleteFindMax :: Map k a -> ((k,a),Map k a)
1216 deleteFindMax t
1217 = case t of
1218 Bin _ k x l Tip -> ((k,x),l)
1219 Bin _ k x l r -> let (km,r') = deleteFindMax r in (km,balance k x l r')
1220 Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip)
1221
1222
1223 {--------------------------------------------------------------------
1224 [balance l x r] balances two trees with value x.
1225 The sizes of the trees should balance after decreasing the
1226 size of one of them. (a rotation).
1227
1228 [delta] is the maximal relative difference between the sizes of
1229 two trees, it corresponds with the [w] in Adams' paper.
1230 [ratio] is the ratio between an outer and inner sibling of the
1231 heavier subtree in an unbalanced setting. It determines
1232 whether a double or single rotation should be performed
1233 to restore balance. It is correspondes with the inverse
1234 of $\alpha$ in Adam's article.
1235
1236 Note that:
1237 - [delta] should be larger than 4.646 with a [ratio] of 2.
1238 - [delta] should be larger than 3.745 with a [ratio] of 1.534.
1239
1240 - A lower [delta] leads to a more 'perfectly' balanced tree.
1241 - A higher [delta] performs less rebalancing.
1242
1243 - Balancing is automatic for random data and a balancing
1244 scheme is only necessary to avoid pathological worst cases.
1245 Almost any choice will do, and in practice, a rather large
1246 [delta] may perform better than smaller one.
1247
1248 Note: in contrast to Adam's paper, we use a ratio of (at least) [2]
1249 to decide whether a single or double rotation is needed. Allthough
1250 he actually proves that this ratio is needed to maintain the
1251 invariants, his implementation uses an invalid ratio of [1].
1252 --------------------------------------------------------------------}
1253 delta,ratio :: Int
1254 delta = 5
1255 ratio = 2
1256
1257 balance :: k -> a -> Map k a -> Map k a -> Map k a
1258 balance k x l r
1259 | sizeL + sizeR <= 1 = Bin sizeX k x l r
1260 | sizeR >= delta*sizeL = rotateL k x l r
1261 | sizeL >= delta*sizeR = rotateR k x l r
1262 | otherwise = Bin sizeX k x l r
1263 where
1264 sizeL = size l
1265 sizeR = size r
1266 sizeX = sizeL + sizeR + 1
1267
1268 -- rotate
1269 rotateL k x l r@(Bin _ _ _ ly ry)
1270 | size ly < ratio*size ry = singleL k x l r
1271 | otherwise = doubleL k x l r
1272
1273 rotateR k x l@(Bin _ _ _ ly ry) r
1274 | size ry < ratio*size ly = singleR k x l r
1275 | otherwise = doubleR k x l r
1276
1277 -- basic rotations
1278 singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3
1279 singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3)
1280
1281 doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4)
1282 doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4)
1283
1284
1285 {--------------------------------------------------------------------
1286 The bin constructor maintains the size of the tree
1287 --------------------------------------------------------------------}
1288 bin :: k -> a -> Map k a -> Map k a -> Map k a
1289 bin k x l r
1290 = Bin (size l + size r + 1) k x l r
1291
1292
1293 {--------------------------------------------------------------------
1294 Eq converts the tree to a list. In a lazy setting, this
1295 actually seems one of the faster methods to compare two trees
1296 and it is certainly the simplest :-)
1297 --------------------------------------------------------------------}
1298 instance (Eq k,Eq a) => Eq (Map k a) where
1299 t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2)
1300
1301 {--------------------------------------------------------------------
1302 Ord
1303 --------------------------------------------------------------------}
1304
1305 instance (Ord k, Ord v) => Ord (Map k v) where
1306 compare m1 m2 = compare (toAscList m1) (toAscList m2)
1307
1308 {--------------------------------------------------------------------
1309 Functor
1310 --------------------------------------------------------------------}
1311 instance Functor (Map k) where
1312 fmap f m = map f m
1313
1314 {--------------------------------------------------------------------
1315 Read
1316 --------------------------------------------------------------------}
1317 instance (Ord k, Read k, Read e) => Read (Map k e) where
1318 #ifdef __GLASGOW_HASKELL__
1319 readPrec = parens $ prec 10 $ do
1320 Ident "fromList" <- lexP
1321 xs <- readPrec
1322 return (fromList xs)
1323
1324 readListPrec = readListPrecDefault
1325 #else
1326 readsPrec p = readParen (p > 10) $ \ r -> do
1327 ("fromList",s) <- lex r
1328 (xs,t) <- reads s
1329 return (fromList xs,t)
1330 #endif
1331
1332 -- parses a pair of things with the syntax a:=b
1333 readPair :: (Read a, Read b) => ReadS (a,b)
1334 readPair s = do (a, ct1) <- reads s
1335 (":=", ct2) <- lex ct1
1336 (b, ct3) <- reads ct2
1337 return ((a,b), ct3)
1338
1339 {--------------------------------------------------------------------
1340 Show
1341 --------------------------------------------------------------------}
1342 instance (Show k, Show a) => Show (Map k a) where
1343 showsPrec d m = showParen (d > 10) $
1344 showString "fromList " . shows (toList m)
1345
1346 showMap :: (Show k,Show a) => [(k,a)] -> ShowS
1347 showMap []
1348 = showString "{}"
1349 showMap (x:xs)
1350 = showChar '{' . showElem x . showTail xs
1351 where
1352 showTail [] = showChar '}'
1353 showTail (x:xs) = showString ", " . showElem x . showTail xs
1354
1355 showElem (k,x) = shows k . showString " := " . shows x
1356
1357
1358 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1359 -- in a compressed, hanging format.
1360 showTree :: (Show k,Show a) => Map k a -> String
1361 showTree m
1362 = showTreeWith showElem True False m
1363 where
1364 showElem k x = show k ++ ":=" ++ show x
1365
1366
1367 {- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows
1368 the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is
1369 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1370 @wide@ is 'True', an extra wide version is shown.
1371
1372 > Map> let t = fromDistinctAscList [(x,()) | x <- [1..5]]
1373 > Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True False t
1374 > (4,())
1375 > +--(2,())
1376 > | +--(1,())
1377 > | +--(3,())
1378 > +--(5,())
1379 >
1380 > Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True True t
1381 > (4,())
1382 > |
1383 > +--(2,())
1384 > | |
1385 > | +--(1,())
1386 > | |
1387 > | +--(3,())
1388 > |
1389 > +--(5,())
1390 >
1391 > Map> putStrLn $ showTreeWith (\k x -> show (k,x)) False True t
1392 > +--(5,())
1393 > |
1394 > (4,())
1395 > |
1396 > | +--(3,())
1397 > | |
1398 > +--(2,())
1399 > |
1400 > +--(1,())
1401
1402 -}
1403 showTreeWith :: (k -> a -> String) -> Bool -> Bool -> Map k a -> String
1404 showTreeWith showelem hang wide t
1405 | hang = (showsTreeHang showelem wide [] t) ""
1406 | otherwise = (showsTree showelem wide [] [] t) ""
1407
1408 showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> ShowS
1409 showsTree showelem wide lbars rbars t
1410 = case t of
1411 Tip -> showsBars lbars . showString "|\n"
1412 Bin sz kx x Tip Tip
1413 -> showsBars lbars . showString (showelem kx x) . showString "\n"
1414 Bin sz kx x l r
1415 -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r .
1416 showWide wide rbars .
1417 showsBars lbars . showString (showelem kx x) . showString "\n" .
1418 showWide wide lbars .
1419 showsTree showelem wide (withEmpty lbars) (withBar lbars) l
1420
1421 showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS
1422 showsTreeHang showelem wide bars t
1423 = case t of
1424 Tip -> showsBars bars . showString "|\n"
1425 Bin sz kx x Tip Tip
1426 -> showsBars bars . showString (showelem kx x) . showString "\n"
1427 Bin sz kx x l r
1428 -> showsBars bars . showString (showelem kx x) . showString "\n" .
1429 showWide wide bars .
1430 showsTreeHang showelem wide (withBar bars) l .
1431 showWide wide bars .
1432 showsTreeHang showelem wide (withEmpty bars) r
1433
1434
1435 showWide wide bars
1436 | wide = showString (concat (reverse bars)) . showString "|\n"
1437 | otherwise = id
1438
1439 showsBars :: [String] -> ShowS
1440 showsBars bars
1441 = case bars of
1442 [] -> id
1443 _ -> showString (concat (reverse (tail bars))) . showString node
1444
1445 node = "+--"
1446 withBar bars = "| ":bars
1447 withEmpty bars = " ":bars
1448
1449 {--------------------------------------------------------------------
1450 Typeable
1451 --------------------------------------------------------------------}
1452
1453 #include "Typeable.h"
1454 INSTANCE_TYPEABLE2(Map,mapTc,"Map")
1455
1456 {--------------------------------------------------------------------
1457 Assertions
1458 --------------------------------------------------------------------}
1459 -- | /O(n)/. Test if the internal map structure is valid.
1460 valid :: Ord k => Map k a -> Bool
1461 valid t
1462 = balanced t && ordered t && validsize t
1463
1464 ordered t
1465 = bounded (const True) (const True) t
1466 where
1467 bounded lo hi t
1468 = case t of
1469 Tip -> True
1470 Bin sz kx x l r -> (lo kx) && (hi kx) && bounded lo (<kx) l && bounded (>kx) hi r
1471
1472 -- | Exported only for "Debug.QuickCheck"
1473 balanced :: Map k a -> Bool
1474 balanced t
1475 = case t of
1476 Tip -> True
1477 Bin sz kx x l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) &&
1478 balanced l && balanced r
1479
1480
1481 validsize t
1482 = (realsize t == Just (size t))
1483 where
1484 realsize t
1485 = case t of
1486 Tip -> Just 0
1487 Bin sz kx x l r -> case (realsize l,realsize r) of
1488 (Just n,Just m) | n+m+1 == sz -> Just sz
1489 other -> Nothing
1490
1491 {--------------------------------------------------------------------
1492 Utilities
1493 --------------------------------------------------------------------}
1494 foldlStrict f z xs
1495 = case xs of
1496 [] -> z
1497 (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1498
1499
1500 {-
1501 {--------------------------------------------------------------------
1502 Testing
1503 --------------------------------------------------------------------}
1504 testTree xs = fromList [(x,"*") | x <- xs]
1505 test1 = testTree [1..20]
1506 test2 = testTree [30,29..10]
1507 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1508
1509 {--------------------------------------------------------------------
1510 QuickCheck
1511 --------------------------------------------------------------------}
1512 qcheck prop
1513 = check config prop
1514 where
1515 config = Config
1516 { configMaxTest = 500
1517 , configMaxFail = 5000
1518 , configSize = \n -> (div n 2 + 3)
1519 , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1520 }
1521
1522
1523 {--------------------------------------------------------------------
1524 Arbitrary, reasonably balanced trees
1525 --------------------------------------------------------------------}
1526 instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where
1527 arbitrary = sized (arbtree 0 maxkey)
1528 where maxkey = 10000
1529
1530 arbtree :: (Enum k,Arbitrary a) => Int -> Int -> Int -> Gen (Map k a)
1531 arbtree lo hi n
1532 | n <= 0 = return Tip
1533 | lo >= hi = return Tip
1534 | otherwise = do{ x <- arbitrary
1535 ; i <- choose (lo,hi)
1536 ; m <- choose (1,30)
1537 ; let (ml,mr) | m==(1::Int)= (1,2)
1538 | m==2 = (2,1)
1539 | m==3 = (1,1)
1540 | otherwise = (2,2)
1541 ; l <- arbtree lo (i-1) (n `div` ml)
1542 ; r <- arbtree (i+1) hi (n `div` mr)
1543 ; return (bin (toEnum i) x l r)
1544 }
1545
1546
1547 {--------------------------------------------------------------------
1548 Valid tree's
1549 --------------------------------------------------------------------}
1550 forValid :: (Show k,Enum k,Show a,Arbitrary a,Testable b) => (Map k a -> b) -> Property
1551 forValid f
1552 = forAll arbitrary $ \t ->
1553 -- classify (balanced t) "balanced" $
1554 classify (size t == 0) "empty" $
1555 classify (size t > 0 && size t <= 10) "small" $
1556 classify (size t > 10 && size t <= 64) "medium" $
1557 classify (size t > 64) "large" $
1558 balanced t ==> f t
1559
1560 forValidIntTree :: Testable a => (Map Int Int -> a) -> Property
1561 forValidIntTree f
1562 = forValid f
1563
1564 forValidUnitTree :: Testable a => (Map Int () -> a) -> Property
1565 forValidUnitTree f
1566 = forValid f
1567
1568
1569 prop_Valid
1570 = forValidUnitTree $ \t -> valid t
1571
1572 {--------------------------------------------------------------------
1573 Single, Insert, Delete
1574 --------------------------------------------------------------------}
1575 prop_Single :: Int -> Int -> Bool
1576 prop_Single k x
1577 = (insert k x empty == singleton k x)
1578
1579 prop_InsertValid :: Int -> Property
1580 prop_InsertValid k
1581 = forValidUnitTree $ \t -> valid (insert k () t)
1582
1583 prop_InsertDelete :: Int -> Map Int () -> Property
1584 prop_InsertDelete k t
1585 = (lookup k t == Nothing) ==> delete k (insert k () t) == t
1586
1587 prop_DeleteValid :: Int -> Property
1588 prop_DeleteValid k
1589 = forValidUnitTree $ \t ->
1590 valid (delete k (insert k () t))
1591
1592 {--------------------------------------------------------------------
1593 Balance
1594 --------------------------------------------------------------------}
1595 prop_Join :: Int -> Property
1596 prop_Join k
1597 = forValidUnitTree $ \t ->
1598 let (l,r) = split k t
1599 in valid (join k () l r)
1600
1601 prop_Merge :: Int -> Property
1602 prop_Merge k
1603 = forValidUnitTree $ \t ->
1604 let (l,r) = split k t
1605 in valid (merge l r)
1606
1607
1608 {--------------------------------------------------------------------
1609 Union
1610 --------------------------------------------------------------------}
1611 prop_UnionValid :: Property
1612 prop_UnionValid
1613 = forValidUnitTree $ \t1 ->
1614 forValidUnitTree $ \t2 ->
1615 valid (union t1 t2)
1616
1617 prop_UnionInsert :: Int -> Int -> Map Int Int -> Bool
1618 prop_UnionInsert k x t
1619 = union (singleton k x) t == insert k x t
1620
1621 prop_UnionAssoc :: Map Int Int -> Map Int Int -> Map Int Int -> Bool
1622 prop_UnionAssoc t1 t2 t3
1623 = union t1 (union t2 t3) == union (union t1 t2) t3
1624
1625 prop_UnionComm :: Map Int Int -> Map Int Int -> Bool
1626 prop_UnionComm t1 t2
1627 = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1628
1629 prop_UnionWithValid
1630 = forValidIntTree $ \t1 ->
1631 forValidIntTree $ \t2 ->
1632 valid (unionWithKey (\k x y -> x+y) t1 t2)
1633
1634 prop_UnionWith :: [(Int,Int)] -> [(Int,Int)] -> Bool
1635 prop_UnionWith xs ys
1636 = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys)))
1637 == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys))
1638
1639 prop_DiffValid
1640 = forValidUnitTree $ \t1 ->
1641 forValidUnitTree $ \t2 ->
1642 valid (difference t1 t2)
1643
1644 prop_Diff :: [(Int,Int)] -> [(Int,Int)] -> Bool
1645 prop_Diff xs ys
1646 = List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1647 == List.sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys)))
1648
1649 prop_IntValid
1650 = forValidUnitTree $ \t1 ->
1651 forValidUnitTree $ \t2 ->
1652 valid (intersection t1 t2)
1653
1654 prop_Int :: [(Int,Int)] -> [(Int,Int)] -> Bool
1655 prop_Int xs ys
1656 = List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1657 == List.sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys)))
1658
1659 {--------------------------------------------------------------------
1660 Lists
1661 --------------------------------------------------------------------}
1662 prop_Ordered
1663 = forAll (choose (5,100)) $ \n ->
1664 let xs = [(x,()) | x <- [0..n::Int]]
1665 in fromAscList xs == fromList xs
1666
1667 prop_List :: [Int] -> Bool
1668 prop_List xs
1669 = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])])
1670 -}