TrieTree and give suggestions

scala/src/main/scala/ch35_tire_tree/TrieTree.scala

@@ -0,0 +1,106 @@
+package ch35_tire_tree
+
+import scala.collection.mutable.ArrayBuffer
+import scala.util.control.Breaks._
+
+class TrieNode(var data: Char, var children: Array[TrieNode], var isLeafNode: Boolean = false) {
+
+  def this(data: Char) = this(data, new Array[TrieNode](26))
+
+}
+
+class TrieTree {
+
+  val root = new TrieNode('/') //root node without any data
+
+  def insert(text: Array[Char]): Unit = {
+    var p = root
+    for (i <- text.indices) {
+      val index = text(i) - 'a'
+      if (p.children(index) == null) {
+        val node = new TrieNode(text(i))
+        p.children(index) = node
+      }
+      p = p.children(index)
+    }
+    //now p is pointing the leaf node
+    p.isLeafNode = true
+  }
+
+  def find(text: Array[Char]): Boolean = {
+    var p = root
+
+    breakable {
+      for (i <- text.indices) {
+        val index = text(i) - 'a'
+        if (p.children(index) == null) {
+          break
+        }
+        p = p.children(index)
+      }
+    }
+
+    p.isLeafNode
+  }
+
+  def prefix(prefix: Array[Char]): Boolean = {
+    var p = root
+    var level = -1
+    breakable {
+      for (i <- prefix.indices) {
+        val index = prefix(i) - 'a'
+        if (p.children(index) == null) {
+          break
+        }
+        p = p.children(index)
+        level = i
+      }
+    }
+
+    //should not reach leaf node and should have at least 1 matching level
+    !p.isLeafNode && level > -1
+  }
+
+  def suggestion(text: Array[Char]): Option[Array[String]] = {
+    var p = root
+    var level = -1
+    breakable {
+      for (i <- text.indices) {
+        val index = text(i) - 'a'
+        if (p.children(index) == null) {
+          break
+        }
+        p = p.children(index)
+        level = i
+      }
+    }
+    //have prefix matching
+    if (!p.isLeafNode && level > -1) {
+      //now the problem becomes print all the children from p
+      Some(_children(p).map(str => "%s%s".format(text.slice(0, text.length - 1).mkString(""), str)).toArray)
+    } else {
+      None
+    }
+
+  }
+
+  def _children(p: TrieNode): ArrayBuffer[String] = {
+    if (p.isLeafNode) {
+      return ArrayBuffer(p.data.toString)
+    }
+
+    //p is not leaf node
+    val arrayBuffer = new ArrayBuffer[String]()
+    for (i <- p.children.indices) {
+      if (p.children(i) != null) {
+        _children(p.children(i)).foreach(subStr =>
+          arrayBuffer.append("%s%s".format(p.data.toString, subStr))
+        )
+      }
+    }
+
+    arrayBuffer
+  }
+
+
+}

scala/src/test/scala/ch35_tire_tree/TrieTreeTest.scala

@@ -0,0 +1,48 @@
+package ch35_tire_tree
+
+import org.scalatest.{FlatSpec, Matchers}
+
+class TrieTreeTest extends FlatSpec with Matchers {
+
+  behavior of "TrieTreeTest"
+
+  it should "insert data into the tree and find it back" in {
+    val texts = Array("how", "hi", "her", "hello", "so", "see")
+    val tree = new TrieTree()
+    for (i <- texts.indices) {
+      tree.insert(texts(i).toCharArray)
+    }
+
+    for (i <- texts.indices) {
+      assert(tree.find(texts(i).toCharArray))
+    }
+  }
+
+  it should "find prefix" in {
+    val texts = Array("how", "hi", "her", "hello", "so", "see")
+    val tree = new TrieTree()
+    for (i <- texts.indices) {
+      tree.insert(texts(i).toCharArray)
+    }
+
+    for (i <- texts.indices) {
+      assert(tree.prefix(texts(i).toCharArray.slice(0, 1)))
+    }
+
+    assert(!tree.prefix("howlo".toCharArray))
+    assert(!tree.prefix("aaa".toCharArray))
+  }
+
+  it should "give suggestion" in {
+    val texts = Array("how", "hi", "her", "hello", "so", "see")
+    val tree = new TrieTree()
+    for (i <- texts.indices) {
+      tree.insert(texts(i).toCharArray)
+    }
+
+    tree.suggestion("he".toCharArray).get should contain("her")
+    tree.suggestion("he".toCharArray).get should contain("hello")
+    tree.suggestion("he".toCharArray).get.size should equal(2)
+  }
+
+}