When I wrote about Scala linked list implementation a couple of years ago, I also did some quick ground work for implementing binary search trees (BST). Occupied by other R&D projects at the time, it was put aside and has since been patiently awaiting its turn to see the light of day. As much of the code is already there, I’m going to put it up in this blog post along with some narrative remarks.
First, we come up with an ADT (algebraic data type). Let’s call it BSTree, starting out with a base trait with generic type A for the data element to be stored inside the tree structure, to be extended by a case class BSBranch as tree branches and a case object BSLeaf as “null” tree nodes. The ADT’s overall structure resembles that of the one used in the linked list implementation described in the old post.
ADT BSTree
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
sealed trait BSTree[+A] { self => // ... override def toString: String = self match { case BSLeaf => "-" case BSBranch(e, c, l, r) => (l, r) match { case (BSLeaf, BSLeaf) => s"TN($e[$c], -, -)" case (BSBranch(le, lc, _, _), BSLeaf) => s"TN($e[$c], $le[$lc], -)" case (BSLeaf, BSBranch(re, rc, _, _)) => s"TN($e[$c], -, $re[$rc])" case (BSBranch(le, lc, _, _), BSBranch(re, rc, _, _)) => s"TN($e[$c], $le[$lc], $re[$rc])" } } } case class BSBranch[+A]( elem: A, count: Int = 1, left: BSTree[A] = BSLeaf, right: BSTree[A] = BSLeaf ) extends BSTree[A] case object BSLeaf extends BSTree[Nothing] |
A few notes:
- Data elements (of generic type
A) are stored only in branches and data of the same value will go to the same branch and be represented with a propercountvalue. - BSTree is covariant, or else
BSLeafcan’t even be defined as a sub-type ofBSTree[Nothing]. - A
toStringmethod is created for simplified string output of a tree instance.
Populating a BSTree
One of the first things we need is a method to insert tree nodes into an existing BSTree. We start expanding the base trait with method insert(). That’s all great for adding a node one at a time, but we also need a way to create a BSTree and populate it from a readily available collection of data elements. It makes sense to delegate such a factory method to the companion object BSTree as its method apply().
|
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 |
sealed trait BSTree[+A] { self => def insert[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSBranch(elem, 1, BSLeaf, BSLeaf) case t @ BSBranch(e, c, l, r) => if (e > elem) t.copy(left = l.insert(elem)) else if (e < elem) t.copy(right = r.insert(elem)) else t.copy(count = c + 1) } } // ... } object BSTree { def apply[A : Ordering](elems: Vector[A]): BSTree[A] = { val ord = implicitly[Ordering[A]] import ord._ elems.foldLeft[BSTree[A]](BSLeaf)(_.insert(_)) } } |
Note that type parameter B for insert() needs to be a supertype of A because Function1 is contravariant over its parameter type. In addition, the context bound “B : Ordering” constrains type B to be capable of being ordered (i.e. compared) which is necessary for traversing a binary search tree.
Testing BSTree.apply():
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
val tree = BSTree(Vector(30, 10, 50, 20, 40, 70, 60, 80, 20, 50, 60)) // tree: BSTree[Int] = TN(30[1], 10[1], 50[2]) /* tree: 30 / \ 10 50 \ / \ 20 40 70 / \ 60 80 */ |
Tree traversal and finding tree nodes
Next, we need methods for tree traversal and search. For brevity, we only include in-order traversal.
|
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 |
def traverseInOrder: Unit = self match { case BSLeaf => () case BSBranch(_, _, l, r) => l.traverseInOrder println(self) r.traverseInOrder } def traverseByLevel: Unit = { def loop(tree: BSTree[A], level: Int): Unit = { if (level > 0) tree match { case BSLeaf => case BSBranch(_, _, l, r) => loop(l, level - 1) loop(r, level - 1) } else print(s"$tree ") } for (l <- 0 until self.height) { loop(self, l) println() } } def find[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, c, l, r) => if (e < elem) r.find(elem) else if (e > elem) l.find(elem) else t } } def height: Int = self match { case BSLeaf => 0 case BSBranch(_, _, l, r) => (l.height max r.height) + 1 } |
Using the tree created above:
|
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 |
/* tree: 30 / \ 10 50 \ / \ 20 40 70 / \ 60 80 */ tree.height // 4 tree.traverseInOrder /* TN(10[1], -, 20[2]) TN(20[2], -, -) TN(30[1], 10[1], 50[2]) TN(40[1], -, -) TN(50[2], 40[1], 70[1]) TN(60[2], -, -) TN(70[1], 60[2], 80[1]) TN(80[1], -, -) */ tree.traverseByLevel /* TN(30[1], 10[1], 50[2]) TN(10[1], -, 20[2]) TN(50[2], 40[1], 70[1]) TN(20[2], -, -) TN(40[1], -, -) TN(70[1], 60[2], 80[1]) TN(60[2], -, -) TN(80[1], -, -) */ val tree50 = tree.find(50) // TN(50[2], 40[1], 70[1]) /* tree50: 50 / \ 40 70 / \ 60 80 */ |
Removing tree nodes
To be able to remove tree nodes that consist of a specific or range of element values, we include also the following few methods in the base trait.
Note that delete() may involve a little shuffling of the tree nodes. Once the tree node to be removed is located, that node may need to be filled with the node having the next-bigger element & count values from its right node (or equivalently, the node having the next-smaller element from its left node).
|
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 |
def delete[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ def leftMost(tree: BSTree[B], prev: BSTree[B]): BSTree[B] = tree match { case BSLeaf => prev case t @ BSBranch(_, _, l, _) => leftMost(l, t) } self match { case BSLeaf => BSLeaf case t @ BSBranch(e, c, l, r) => if (e < elem) t.copy(right = r.delete(elem)) else if (e > elem) t.copy(left = l.delete(elem)) else { if (l == BSLeaf) r else if (r == BSLeaf) l else { val nextBigger = leftMost(r, r) nextBigger match { case BSBranch(enb, cnb, _, _) => t.copy(elem = enb, count = cnb, right = r.delete(enb)) } } } } } def trim[B >: A : Ordering](lower: B, upper: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, _, l, r) => val tt = t.copy(left = l.trim(lower, upper), right = r.trim(lower, upper)) if (e < lower) tt match { case BSBranch(_, _, _, r) => r } else if (e > upper) tt match { case BSBranch(_, _, l, _) => l } else tt } } def cutOut[B >: A : Ordering](lower: B, upper: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, _, l, r) => val tt = t.copy(left = l.cutOut(lower, upper), right = r.cutOut(lower, upper)) if (e >= lower && e <= upper) tt match { case BSBranch(e, _, _, _) => tt.delete(e) } else tt } } |
Method trim() removes tree nodes with element values below or above the provided range. Meanwhile, method cutOut() does the opposite by cutting out tree nodes with values within the given range. It involves slightly more work than trim(), requiring the use of delete() for individual tree nodes.
Example:
|
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 |
val treeD50 = tree.delete(50) // TN(30[1], 10[1], 60[2]) /* treeD50: 30 / \ 10 60 \ / \ 20 40 70 \ 80 */ val treeT15to65 = tree.trim(15, 65) // TN(30[1], 20[2], 50[2]) /* treeT15to65: 30 / \ 20 50 / \ 40 60 */ val treeC25to55 = tree.trimInRange(25, 55) // TN(60[1], 10[1], 70[1]) /* treeC25to55: 60 / \ 10 70 \ \ 20 80 */ |
Rebuilding a binary search tree
A highly unbalanced binary search tree beats the purpose of using such a data structure. One of the most straight forward ways to rebuild a binary search tree is to “unpack” the individual tree nodes of the existing tree by traversing in-order into a list (e.g. a Vector or List) of elements, followed by reconstructing a new tree with nodes being assigned elements from recursively half-ing the in-order node list.
|
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 |
def rebuild: BSTree[A] = { def loop(vs: Vector[(A, Int)], lower: Int, upper: Int): BSTree[A] = { if (upper >= lower) { val m = (lower + upper) / 2 val (e, c) = vs(m) BSBranch(e, c, loop(vs, lower, m-1), loop(vs, m+1, upper)) } else BSLeaf } self match { case BSLeaf => BSLeaf case _ => val vs = self.toElemVector loop(vs, 0, vs.size - 1) } } def isBalanced: Boolean = self match { case BSLeaf => true case BSBranch(_, _, l, r) => if (math.abs(l.height - r.height) > 1) false else l.isBalanced && r.isBalanced } def toElemVector: Vector[(A, Int)] = self match { // in-order case BSLeaf => Vector.empty[(A, Int)] case BSBranch(e, c, l, r) => (l.toElemVector :+ (e, c)) ++ r.toElemVector } def toVector: Vector[BSTree[A]] = self match { // in-order case BSLeaf => Vector.empty[BSTree[A]] case BSBranch(_, _, l, r) => (l.toVector :+ self) ++ r.toVector } |
Example:
|
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 |
val unbalancedTree = BSTree(Vector(20, 10, 30, 30, 40, 50, 60, 60, 50)) // unbalancedTree: BSTree[Int] = TN(20[1], 10[1], 30[2]) /* unbalancedTree: 20 / \ 10 30 \ 40 \ 60 / 50 */ unbalancedTree.isBalanced // false val rebuiltTree = unbalancedTree.rebuild // TN(30[2], 10[1], 50[2]) /* rebuiltTree: 30 / \ 10 50 \ / \ 20 40 60 */ rebuiltTree.isBalanced // true |
Thoughts on the ADT
An alternative to how the ADT is designed is to have the class fields and methods declared in the BSTree base trait with specific implementations reside within subclasses BSBranch and BSLeaf, thus eliminating the need of the boiler-plate pattern matching for the subclasses. There is also the benefit of making class fields like left & right referenceable from the base trait, though they would need to be wrapped in Options with value None for BSLeaf.
As can be seen with the existing ADT, an advantage is having all the binary tree functions defined within the base trait once and for all. If there is the need for having left and right referenceable from the BSTree base trait, one can define something like below within the trait.
|
1 2 3 4 5 6 7 8 9 |
def leftIfAny: Option[BSTree[A]] = self match { case BSLeaf => None case BSBranch(_, _, l, _) => Some(l) } def rightIfAny: Option[BSTree[A]] = self match { case BSLeaf => None case BSBranch(_, _, _, r) => Some(r) } |
Example:
|
1 2 3 4 5 |
tree.leftIfAny // Some(TN(10[1], -, 20[2])) tree.rightIfAny.flatMap(_.rightIfAny) // Some(TN(70[1], 60[2], 80[1])) |
Then there is also non-idiomatic approach of using mutable class fields in a single tree class commonly seen in Java implementation, like below:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
object SimpleBST { class TreeNode[A](_elem: A, _left: TreeNode[A] = null, _right: TreeNode[A] = null) { var elem: A = _elem var count: Int = 1 var left: TreeNode[A] = _left var right: TreeNode[A] = _right override def toString: String = if (this == null) "null" else (left, right) match { case (null, null) => s"TN($elem[$count], -, -)" case (l, null) => s"TN($elem[$count], ${l.elem}[${l.count}], -)" case (null, r) => s"TN($elem[$count], -, ${r.elem}[${r.count}])" case (l, r) => s"TN($elem[$count], ${l.elem}[${l.count}], ${r.elem}[${r.count}])" } } // Methods for node addition, deletion, etc ... } |
Addendum: Complete source code of the BSTree ADT
|
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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
sealed trait BSTree[+A] { self => def traverseInOrder: Unit = self match { case BSLeaf => () case BSBranch(_, _, l, r) => l.traverseInOrder println(self) r.traverseInOrder } def traverseByLevel: Unit = { def loop(tree: BSTree[A], level: Int): Unit = { if (level > 0) tree match { case BSLeaf => case BSBranch(_, _, l, r) => loop(l, level - 1) loop(r, level - 1) } else print(s"$tree ") } for (l <- 0 until self.height) { loop(self, l) println() } } def find[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, c, l, r) => if (e < elem) r.find(elem) else if (e > elem) l.find(elem) else t } } def height: Int = self match { case BSLeaf => 0 case BSBranch(_, _, l, r) => (l.height max r.height) + 1 } def insert[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSBranch(elem, 1, BSLeaf, BSLeaf) case t @ BSBranch(e, c, l, r) => if (e > elem) t.copy(left = l.insert(elem)) else if (e < elem) t.copy(right = r.insert(elem)) else t.copy(count = c + 1) } } def delete[B >: A : Ordering](elem: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ def leftMost(tree: BSTree[B], prev: BSTree[B]): BSTree[B] = tree match { case BSLeaf => prev case t @ BSBranch(_, _, l, _) => leftMost(l, t) } self match { case BSLeaf => BSLeaf case t @ BSBranch(e, c, l, r) => if (e < elem) t.copy(right = r.delete(elem)) else if (e > elem) t.copy(left = l.delete(elem)) else { if (l == BSLeaf) r else if (r == BSLeaf) l else { val nextBigger = leftMost(r, r) nextBigger match { case BSBranch(enb, cnb, _, _) => t.copy(elem = enb, count = cnb, right = r.delete(enb)) } } } } } def trim[B >: A : Ordering](lower: B, upper: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, _, l, r) => val tt = t.copy(left = l.trim(lower, upper), right = r.trim(lower, upper)) if (e < lower) tt match { case BSBranch(_, _, _, r) => r } else if (e > upper) tt match { case BSBranch(_, _, l, _) => l } else tt } } def cutOut[B >: A : Ordering](lower: B, upper: B): BSTree[B] = { val ord = implicitly[Ordering[B]] import ord._ self match { case BSLeaf => BSLeaf case t @ BSBranch(e, _, l, r) => val tt = t.copy(left = l.cutOut(lower, upper), right = r.cutOut(lower, upper)) if (e >= lower && e <= upper) tt match { case BSBranch(e, _, _, _) => tt.delete(e) } else tt } } def rebuild(): BSTree[A] = { def loop(vs: Vector[A], lower: Int, upper: Int): BSTree[A] = { if (upper >= lower) { val m = (lower + upper) / 2 val em = vs(m) self match { case BSBranch(e, c, _, _) if e == em => BSBranch(em, c+1, loop(vs, lower, m-1), loop(vs, m+1, upper)) case _ => BSBranch(em, 1, loop(vs, lower, m-1), loop(vs, m+1, upper)) } } else BSLeaf } self match { case BSLeaf => BSLeaf case _ => val vs = self.toElemVector loop(vs, 0, vs.size - 1) } } def isBalanced: Boolean = self match { case BSLeaf => true case BSBranch(_, _, l, r) => if (math.abs(l.height - r.height) > 1) false else l.isBalanced && r.isBalanced } def toElemVector: Vector[A] = self match { // in-order case BSLeaf => Vector.empty[A] case BSBranch(e, _, l, r) => (l.toElemVector :+ e) ++ r.toElemVector } def toVector: Vector[BSTree[A]] = self match { // in-order case BSLeaf => Vector.empty[BSTree[A]] case BSBranch(_, _, l, r) => (l.toVector :+ self) ++ r.toVector } override def toString: String = self match { case BSLeaf => "-" case BSBranch(e, c, l, r) => (l, r) match { case (BSLeaf, BSLeaf) => s"TN($e[$c], -, -)" case (BSBranch(le, lc, _, _), BSLeaf) => s"TN($e[$c], $le[$lc], -)" case (BSLeaf, BSBranch(re, rc, _, _)) => s"TN($e[$c], -, $re[$rc])" case (BSBranch(le, lc, _, _), BSBranch(re, rc, _, _)) => s"TN($e[$c], $le[$lc], $re[$rc])" } } } object BSTree { def apply[A : Ordering](elems: Vector[A]): BSTree[A] = { val ord = implicitly[Ordering[A]] import ord._ elems.foldLeft[BSTree[A]](BSLeaf)(_.insert(_)) } } case class BSBranch[A]( elem: A, count: Int = 1, left: BSTree[A] = BSLeaf, right: BSTree[A] = BSLeaf ) extends BSTree[A] case object BSLeaf extends BSTree[Nothing] |