class Color case object Red extends Color case object Black extends Co

 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
class Color
case object Red extends Color
case object Black extends Color
abstract class RBTree[K, V](implicit cmp: K => Ordered[K]) {
protected def blacken(n: RBTree[K, V]): RBTree[K, V] = {
n match {
case Leaf() => n
case Branch(_, l, k, v, r) => Branch(Black, l, k, v, r)
}
}
protected def balance(color: Color, left: RBTree[K, V], key: K, value: Option[V], right: RBTree[K, V]): RBTree[K, V] = {
(color, left, key, value, right) match {
case (Black, Branch(Red, Branch(Red, a, xK, xV, b), yK, yV, c), zK, zV, d) => Branch(Red, Branch(Black, a, xK, xV, b), yK, yV, Branch(Black, c, zK, zV, d))
case (Black, Branch(Red, a, xK, xV, Branch(Red, b, yK, yV, c)), zK, zV, d) => Branch(Red, Branch(Black, a, xK, xV, b), yK, yV, Branch(Black, c, zK, zV, d))
case (Black, a, xK, xV, Branch(Red, Branch(Red, b, yK, yV, c), zK, zV, d)) => Branch(Red, Branch(Black, a, xK, xV, b), yK, yV, Branch(Black, c, zK, zV, d))
case (Black, a, xK, xV, Branch(Red, b, yK, yV, Branch(Red, c, zK, zV, d))) => Branch(Red, Branch(Black, a, xK, xV, b), yK, yV, Branch(Black, c, zK, zV, d))
case (c, a, xK, xV, b) => Branch(c, a, xK, xV, b)
}
}
def helperFroActions(key: K, f: (K, Option[V]) => Option[V]): RBTree[K, V]
def doAction(key: K, f: (K, Option[V]) => Option[V]): RBTree[K, V] =
blacken(helperFroActions(key, f))
def get(key: K): Option[V]
def insert(key: K, v: V) = doAction(key, (_, _) => Some(v))
def remove(key: K) = doAction(key, (_, _) => None)
}
private case class Leaf[K, V](implicit cmp: K => Ordered[K]) extends RBTree[K, V] {
def get(key: K): Option[V] = None
def helperFroActions(key: K, f: (K, Option[V]) => Option[V]): RBTree[K, V] = {
Branch(Red, this, key, f(key, None), this)
}
}
private case class Branch[K, V](color: Color, left: RBTree[K, V], key: K, value: Option[V], right: RBTree[K, V])
(implicit cmp: K => Ordered[K]) extends RBTree[K, V] {
def get(key: K): Option[V] = {
if (key < this.key) left.get(key)
else
if (key > this.key) right.get(key)
else
value
}
def helperFroActions(key: K, f: (K, Option[V]) => Option[V]): RBTree[K, V] = {
if (key < this.key)
balance(color, left.helperFroActions(key, f), this.key, this.value, right)
else
if (key == this.key)
Branch(color, left, key, f(this.key, this.value), right)
else
balance(color, left, this.key, this.value, right.helperFroActions(key, f))
}
}
object RBTree {
def empty[K <: Ordered[K], V]: RBTree[K, V] = Leaf()((k: K) => k)
def main(args: Array[String]) {
val testTree: RBTree[Int, String] = Branch(Black, Leaf[Int, String], 11, Option.apply[String]("dd"), Leaf[Int, String])
println(testTree.insert(1, "11").insert(2, "22").insert(3, "33").insert(4, "44").remove(2))
}
}