flying3615
1/11/2017 - 3:28 AM

Pretty printed TreeNode with FunSuite test plus Java version

Pretty printed TreeNode with FunSuite test plus Java version

class TreeNode<T> {
    private T data;
    private List<TreeNode<T>> children;

    TreeNode(T data, List<TreeNode<T>> children) {
        this.data = data;
        this.children = children;
    }

    public T getData() {
        return data;
    }

    public List<TreeNode<T>> getChildren() {
        return children;
    }

}


public class ShowTree {

    private static String asciiDisplay(TreeNode<String> root) {
        int indent = 0;
        StringBuilder sb = new StringBuilder();
        asciiTail(root, indent, sb);
        return sb.toString();
    }

    private static void asciiTail(TreeNode<String> node, int indent, StringBuilder sb) {
        sb.append(getIndentString(indent)+"+-"+node.getData()+"\n");
        for (TreeNode<String> t : node.getChildren()) {
            if (t.getChildren() == null) {
                sb.append(getIndentString(indent+1)+"+-"+t.getData()+"\n");
            } else {
                asciiTail(t, indent + 1, sb);
            }
        }

    }

    private static String getIndentString(int indent) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < indent; i++) {
            if(i==0||i>1){
                sb.append("  ");
            }else{
                sb.append("|  ");
            }
        }
        return sb.toString();
    }


    public static void main(String[] args) {
        TreeNode<String> tn1 = new TreeNode<>("Root",
                Arrays.asList(new TreeNode("level1-2", null),
                        new TreeNode("level1-2", null),
                        new TreeNode("level1-3", null)
                ));

        TreeNode<String> tn2 = new TreeNode<>("Root",
                Arrays.asList(new TreeNode("level1-2",
                                Arrays.asList(new TreeNode<>("level2-1",
                                        Arrays.asList(new TreeNode<>("level3-1",
                                                null))))),
                        new TreeNode("level1-2", null),
                        new TreeNode("level1-3", null)
                ));

        System.out.println(asciiDisplay(tn1));
        System.out.println(asciiDisplay(tn2));
    }

}
package scalaInterview

import scala.collection.mutable.ArrayBuffer

case class TreeNode[T](data: T, children: Seq[TreeNode[T]] = Nil)

object ShowTree extends App {

  def asciiDisplay(root: TreeNode[String]): Seq[String] = {

    val sb:ArrayBuffer[String] = new ArrayBuffer[String]()

    def asciiTail(node:TreeNode[String], indent:Int, sb:ArrayBuffer[String]):Unit = {
      sb+=getIndentString(indent)+"+-"+node.data
      node.children.foreach(t=>{
        if(t.children==Nil){
          sb+=(getIndentString(indent+1)+"+-"+t.data)
        }else{
          asciiTail(t,indent+1,sb)
        }
      })
    }
    asciiTail(root,0, sb)
    sb
  }

  def getIndentString(indent:Int) = {
    val sb:StringBuilder = new StringBuilder();
    for(i<-(0 until indent)){
      i match {
        case 0 => sb.append("  ")
        case x if x>1 => sb.append("  ")
        case _ => sb.append("|  ")
      }
    }
    sb.toString()
  }


  asciiDisplay(TreeNode("Root",
    children = List(TreeNode("level1-1"),
      TreeNode("level1-2"),
      TreeNode("level1-3")))).foreach(println)

  println("------------------------")

  asciiDisplay(TreeNode("Root",
    children = List(
      TreeNode("level1-1", children = TreeNode("level2-1", children = TreeNode("level3-1") :: Nil) :: Nil),
      TreeNode("level1-2"),
      TreeNode("level1-3")))).foreach(println)


}


class ShowTree$Test extends FunSuite {

  test("test level 1") {
    val t1 = TreeNode("Root",
      children = List(TreeNode("level1-1"),
        TreeNode("level1-2"),
        TreeNode("level1-3")))
    val result1 = ShowTree.asciiDisplay(t1)

    val expectedResult1 = """+-Root
                           |  +-level1-1
                           |  +-level1-2
                           |  +-level1-3""".stripMargin.lines.toArray

    for(i<-(0 until expectedResult1.size)){
      println(result1(i))
      assert(result1(i) === expectedResult1(i))
    }

  }

  test("test multi-level") {
    val t2 = TreeNode("Root",
      children = List(
        TreeNode("level1-1", children = TreeNode("level2-1", children = TreeNode("level3-1") :: Nil) :: Nil),
        TreeNode("level1-2"),
        TreeNode("level1-3")))
    val result2 = ShowTree.asciiDisplay(t2)
    val expectedResult2 = """+-Root
                            |  +-level1-1
                            |  |  +-level2-1
                            |  |    +-level3-1
                            |  +-level1-2
                            |  +-level1-3""".stripMargin.lines.toArray

    for(i<-(0 until expectedResult2.size)){
      println(result2(i))
      assert(result2(i) === expectedResult2(i))
    }
  }

}