August | 2022 | Genuine Blog

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

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 proper count value.
BSTree is covariant, or else BSLeaf can’t even be defined as a sub-type of BSTree[Nothing].
A toString method 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().

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():

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.

  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:

/*
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).

  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:

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.

  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:

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.

  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:

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:

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

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]

Genuine Blog

A Tech Blog by Leo Cheung

Monthly Archives: August 2022

Scala Binary Search Tree

ADT BSTree

Populating a BSTree

Tree traversal and finding tree nodes

Removing tree nodes

Rebuilding a binary search tree

Thoughts on the ADT

Addendum: Complete source code of the BSTree ADT