1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

----------------------------------------------------------------------------
--- Library with an implementation of sets as red-black trees.
---
--- All the operations on sets are generic, i.e., one has to provide
--- an explicit order predicate `(<)` (less-than) on elements.
---
--- @author Johannes Koj, Michael Hanus, Bernd Brassel
--- @version March 2013
--- @category algorithm
----------------------------------------------------------------------------

module SetRBT where

import qualified RedBlackTree as RBT
import           Maybe               (isJust)

type SetRBT a = RBT.RedBlackTree a

--- Returns an empty set, i.e., an empty red-black tree
--- augmented with an order predicate.
emptySetRBT :: Eq a => (a -> a -> Bool) -> SetRBT a
emptySetRBT = RBT.empty (==) (==)

--- Test for an empty set.
isEmptySetRBT :: SetRBT _ -> Bool
isEmptySetRBT = RBT.isEmpty

--- Returns true if an element is contained in a (red-black tree) set.
--- @param e - an element to be checked for containment
--- @param s - a set (represented as a red-black tree)
--- @return True if e is contained in s
elemRBT :: a -> SetRBT a -> Bool
elemRBT e = isJust . (RBT.lookup e)

--- Inserts an element into a set if it is not already there.
insertRBT :: a -> SetRBT a -> SetRBT a
insertRBT = RBT.update

--- Inserts an element into a multiset.
--- Thus, the same element can have several occurrences in the multiset.
insertMultiRBT :: Eq a => a -> SetRBT a -> SetRBT a
insertMultiRBT e = RBT.setInsertEquivalence (==)
                 . RBT.update e
                 . RBT.setInsertEquivalence (\ _ _ -> False)

--- delete an element from a set.
--- Deletes only a single element from a multi set
deleteRBT :: a -> SetRBT a -> SetRBT a
deleteRBT = RBT.delete

--- Transforms a (red-black tree) set into an ordered list of its elements.
setRBT2list :: SetRBT a -> [a]
setRBT2list = RBT.tree2list

--- Computes the union of two (red-black tree) sets.
--- This is done by inserting all elements of the first set into the
--- second set.
unionRBT :: SetRBT a -> SetRBT a -> SetRBT a
unionRBT s1 s2 = foldr insertRBT s2 (setRBT2list s1)

--- Computes the intersection of two (red-black tree) sets.
--- This is done by inserting all elements of the first set
--- contained in the second set into a new set, which order
--- is taken from the first set.
intersectRBT :: SetRBT a -> SetRBT a -> SetRBT a
intersectRBT s1 s2 = foldr insertRBT (RBT.newTreeLike s1)
                     (filter (`elemRBT` s2) (setRBT2list s1))

--- Generic sort based on insertion into red-black trees.
--- The first argument is the order for the elements.
sortRBT  :: Eq a => (a -> a -> Bool) -> [a] -> [a]
sortRBT = RBT.sortBy