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