Mutable Tree

Tree.scala

package tree.Tree

class Tree(rootValue: Int) {

  class Node(var parent: Option[Node], var children: Set[Node], var data: Int) {
  }

  var root : Node = new Node(None, Set(), rootValue)

  def addLeaf(parent: Node, data: Int): Node = {
      val n = new Node(Some(parent), Set(), data)
      parent.children = parent.children + n
      n
  }

  // This is not really necessary
  def traverse(node: Node): List[Int] = {
    val ints: List[Int] = node.children.flatMap(traverse).toList
    node.data :: ints
  }

  def moveSubtree(node: Node, newParent: Node): Unit = {
    node.parent.get.children = node.parent.get.children - node
    node.parent = Some(newParent)
    node.parent.get.children += node
  }

  def ancestors(node: Node): List[Node] = node.parent match {
    case None => List()
    case Some(parent) => parent :: ancestors(parent)
  }

  def descendants(node: Node): List[Node] = {
    val nodes: List[Node] = node.children.flatMap(descendants).toList
    node :: nodes
  }

  def delete(node: Node): Unit = {
    if (node.children.size > 0) {
      throw new Exception("has children")
    }
    else {
      node.parent.get.children -= node
      node.parent = None
    }
  }
}

TreeTest.scala

package tree.Tree


import org.scalatest.FunSuite
import tree.Tree.Tree

class MySuite extends FunSuite {

  test("potato") {
    val t = new Tree(0)
    val node = t.addLeaf(t.root, 1)
    t.addLeaf(node, 5)

    assert(t.traverse(t.root) == List(0, 1, 5))
  }

  test("move subtree") {
    val t = new Tree(0)
    val adminnode = t.addLeaf(t.root, 1)
    val node5 = t.addLeaf(adminnode, 5)
    val node6 = t.addLeaf(adminnode, 6)
    t.addLeaf(node6, 61)

    assert(t.traverse(t.root).toSet == List(0, 1, 5, 6, 61).toSet)


    // (root
    //   (admin
    //     (5)
    //     (6 61)))
    assert(node5.parent.get == adminnode)
    assert(node6.parent.get == adminnode)
    assert(adminnode.children.contains(node5))
    assert(adminnode.children.contains(node6))
    t.moveSubtree(node5, node6)
    // (root
    //   (admin
    //     (5
    //       (6 61))))
    assert(node6.parent.get == adminnode)
    assert(node5.parent.get == node6)
    assert(!adminnode.children.contains(node5))
    assert(adminnode.children.contains(node6))
    assert(node6.children.contains(node5))
  }

  test("ancestors") {
    val t = new Tree(0)
    val adminnode = t.addLeaf(t.root, 1)
    val node5 = t.addLeaf(adminnode, 5)
    val node6 = t.addLeaf(adminnode, 6)
    val node61 = t.addLeaf(node6, 61)
    assert(t.ancestors(node61).map(a => a.data) == List(6, 1, 0))
  }

  test("descendants") {
    val t = new Tree(0)
    val adminnode = t.addLeaf(t.root, 1)
    val node5 = t.addLeaf(adminnode, 5)
    val node6 = t.addLeaf(adminnode, 6)
    val node61 = t.addLeaf(node6, 61)
    assert(t.descendants(t.root).map(a => a.data) == List(0, 1, 5, 6, 61))
  }

  test("delete node") {
    val t = new Tree(0)
    val adminnode = t.addLeaf(t.root, 1)
    val node5 = t.addLeaf(adminnode, 5)
    t.delete(node5)
    assert(adminnode.children.size == 0)
  }

  test("delete node with children") {
    val t = new Tree(0)
    val adminnode = t.addLeaf(t.root, 1)
    val node5 = t.addLeaf(adminnode, 5)
    assertThrows[Exception](
      t.delete(adminnode)
    )
  }
}