-
Notifications
You must be signed in to change notification settings - Fork 2
/
RedBlackTree.scala
103 lines (83 loc) · 2.95 KB
/
RedBlackTree.scala
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package dataStructures
import dataStructures.miscellaneous.DataStructuresExceptionMessages.EmptyRBTreeException
/**
* Colors of a tree
*/
abstract sealed class Color
case object Red extends Color
case object Black extends Color
abstract sealed class RedBlackTree[+A] {
/**
* Color of a tree
*/
def color: Color
/**
* Current value
*/
def value: Int
/**
* The left branch of a tree
*/
def left: RedBlackTree[Int]
/**
* The right branch of a tree
*/
def right: RedBlackTree[Int]
/**
* Checks whether this tree is empty or not
*/
def isEmpty: Boolean
/**
* Adds an element into a tree
*/
def add(x: Int): RedBlackTree[Int] = {
def balancedAdd(t: RedBlackTree[Int]): RedBlackTree[Int] =
if (t.isEmpty) RedBlackTree.make(Red, x)
else if (x < t.value) balanceLeft(t.color, t.value, balancedAdd(t.left), t.right)
else if (x > t.value) balanceRight(t.color, t.value, t.left, balancedAdd(t.right))
else t
def balanceLeft(c: Color, x: Int, l: RedBlackTree[Int], r: RedBlackTree[Int]) = (c, l, r) match {
case (Black, RBBranch(Red, y, RBBranch(Red, z, a, b), c), d) =>
RedBlackTree.make(Red, y, RedBlackTree.make(Black, z, a, b), RedBlackTree.make(Black, x, c, d))
case (Black, RBBranch(Red, z, a, RBBranch(Red, y, b, c)), d) =>
RedBlackTree.make(Red, y, RedBlackTree.make(Black, z, a, b), RedBlackTree.make(Black, x, c, d))
case _ => RedBlackTree.make(c, x, l, r)
}
def balanceRight(c: Color, x: Int, l: RedBlackTree[Int], r: RedBlackTree[Int]) = (c, l, r) match {
case (Black, a, RBBranch(Red, y, b, RBBranch(Red, z, c, d))) =>
RedBlackTree.make(Red, y, RedBlackTree.make(Black, x, a, b), RedBlackTree.make(Black, z, c, d))
case (Black, a, RBBranch(Red, z, RBBranch(Red, y, b, c), d)) =>
RedBlackTree.make(Red, y, RedBlackTree.make(Black, x, a, b), RedBlackTree.make(Black, z, c, d))
case _ => RedBlackTree.make(c, x, l, r)
}
def blacken(t: RedBlackTree[Int]) = RedBlackTree.make(Black, t.value, t.left, t.right)
blacken(balancedAdd(this))
}
def height: Int =
if (isEmpty) 0
else math.max(left.height, right.height) + 1
}
case class RBBranch[Int](color: Color,
value: Int,
left: RedBlackTree[Int],
right: RedBlackTree[Int]) extends RedBlackTree[Int] {
def isEmpty = false
}
case object Leaf extends RedBlackTree[Nothing] {
def color: Color = Black
def value: Nothing = EmptyRBTreeException()
def left: RedBlackTree[Nothing] = EmptyRBTreeException()
def right: RedBlackTree[Nothing] = EmptyRBTreeException()
def isEmpty = true
}
object RedBlackTree {
/**
* Returns an empty red-black tree
*/
def empty[Int]: RedBlackTree[Int] = Leaf
/**
* Creates a branch
*/
def make(c: Color, x: Int, l: RedBlackTree[Int] = Leaf, r: RedBlackTree[Int] = Leaf): RedBlackTree[Int] =
RBBranch(c, x, l, r)
}