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