Filter VIP nodes of a tree. (non tail-recursive yet)

VipTreeApp.scala

package example

import collection.mutable.ArrayBuffer
import collection.mutable
import annotation.tailrec

/**
 * @author Alex Siman <[email protected]>
 * Date: 2013-09-12
 */
object VipTreeApp extends App {

  def buildTree: List[Node] = {
    val n1 = Node("1", false, None)

    val n1_1 = Node("1", true, Some(n1))
    val n1_2 = Node("2", false, Some(n1))

    val n1_1_1 = Node("1", false, Some(n1_1))
    val n1_1_2 = Node("2", true, Some(n1_1))
    val n1_2_1 = Node("1", true, Some(n1_2))

    val n1_2_1_1 = Node("1", true, Some(n1_2_1))
    val n1_2_1_2 = Node("2", false, Some(n1_2_1))
    val n1_2_1_3 = Node("3", false, Some(n1_2_1))
    val n1_2_1_4 = Node("4", true, Some(n1_2_1))

    val n1_2_1_3_1 = Node("1", false, Some(n1_2_1_3))
    val n1_2_1_3_2 = Node("2", true, Some(n1_2_1_3))
    val n1_2_1_3_3 = Node("3", true, Some(n1_2_1_3))
    val n1_2_1_3_4 = Node("4", false, Some(n1_2_1_3))

    // Root can be represented as list of nodes.
    val tree = List(n1)
    tree
  }

//  @tailrec
  def printTree(nodes: List[Node]) {
    def printTreeImpl(nodes: List[Node], level: Int = 1) {
      if (nodes.nonEmpty) {
        val branchLine = (1 to level * 2).foldLeft("")((acc, i) => acc + "-")
        nodes.foreach { node =>
          println(branchLine + node.fullName + (if (node.isVip) " *" else ""))
          printTreeImpl(node.children, level + 1)
        }
      }
    }
    printTreeImpl(nodes)
    println()
  }

  // TODO: Make this @tailrec
//  @tailrec
  def filterVipNodes(nodes: List[Node]): List[Node] = {
    val res = ArrayBuffer[Node]()
    nodes.foreach(n => {
      val newCs = n.children match {
        case Nil => Nil
        case cs => filterVipNodes(cs)
      }
      if (n.isVip) {
        n.children = newCs
        res += n
      } else if (newCs.nonEmpty) {
        res ++= newCs
      }
    })
    res.toList
  }

  val tree = buildTree
  printTree(tree)

  val vips = filterVipNodes(tree)
  printTree(vips)
}

class Node(
  name: String,
  var isVip: Boolean = false,
  var parent: Option[Node] = None
) {

  val fullName: String = parent.map(_.fullName + ".").getOrElse("") + name

  private var _children: mutable.Buffer[Node] = ArrayBuffer()

  parent.map(_._children += this)

  def children: List[Node] = _children.toList

  def children_=(newCs: List[Node]) {
    _children = newCs.toBuffer
  }
}

object Node {

  def apply(
    name: String,
    isVip: Boolean = false,
    parent: Option[Node] = None
  ): Node = new Node(
    name = name,
    isVip = isVip,
    parent = parent
  )
}

Cool stuff. Very thanks for it. I will use it in my program