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