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