import scala.collection.immutable.HashMap
class MinimizedTrie[T](leaves: HashMap[String, MinimizedTrie[T]], endPoint: Boolean, childrenSize: Int, data: Option[T]) {
def this(endPoint: Boolean, data: Option[T]) = this(HashMap[String, MinimizedTrie[T]](), endPoint, 0, data)
private def leavesNumber() = leaves.size
private def getSingleLeaf() = leaves.last
private def countWords(leaves: HashMap[String, MinimizedTrie[T]]) = leaves.foldLeft(0)(_ + _._2.size())
private def findMaxPrefix(word: String): String = {
for (edge <- leaves.keys) {
val prefix = (word, edge).zipped.takeWhile(Function.tupled(_ == _)).unzip._1.mkString
if (!prefix.isEmpty) {
return prefix
}
}
""
}
def size(): Int = {
if (endPoint) {
return childrenSize + 1
} else {
return childrenSize
}
}
def add(word: String, value: T): MinimizedTrie[T] = {
if (word.isEmpty) {
return new MinimizedTrie[T](leaves, true, childrenSize, Option[T](value))
}
val maxPrefix = findMaxPrefix(word)
if (maxPrefix.isEmpty) {
val newLeaf = word -> new MinimizedTrie[T](true, Option[T](value))
return new MinimizedTrie[T](leaves + newLeaf, endPoint, childrenSize + 1, data)
}
val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
val edge = keyVal._1
val trie = keyVal._2
if (edge == maxPrefix) {
val newLeaf = edge -> trie.add(word.substring(edge.length), value)
val newLeaves = leaves + newLeaf
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
} else {
val end = edge.substring(maxPrefix.length)
val middleVertex = new MinimizedTrie[T](HashMap(end -> trie), false, trie.size(), None)
val newLeaf = maxPrefix -> middleVertex.add(word.substring(maxPrefix.length), value)
val newLeaves = (leaves - edge) + newLeaf
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
}
}
def remove(word: String): MinimizedTrie[T] = {
//this.print()
//println()
if (word.isEmpty) {
return new MinimizedTrie(leaves, false, childrenSize, None)
}
val maxPrefix = findMaxPrefix(word)
if (maxPrefix.isEmpty) {
return this
}
val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
val edge = keyVal._1
val trie = keyVal._2
if (edge == maxPrefix) {
val newLeaf = edge -> trie.remove(word.substring(edge.length))
if (newLeaf._2.size() == 0) {
val newLeaves = leaves - edge
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
} else {
newLeaf._2.leavesNumber() match {
case 1 =>
val nextLeaf = newLeaf._2.getSingleLeaf()
if (newLeaf._2.size() == nextLeaf._2.size()) {
val leaf = (edge + nextLeaf._1) -> nextLeaf._2
val newLeaves = (leaves - edge) + leaf
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
} else {
val newLeaves = leaves + newLeaf
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
}
case _ =>
val newLeaves = leaves + newLeaf
val size = countWords(newLeaves)
return new MinimizedTrie[T](newLeaves, endPoint, size, data)
}
}
} else {
return this
}
}
def get(word: String): Option[T] = {
if (word.isEmpty) {
if (endPoint) {
return data
} else {
return None
}
}
val maxPrefix = findMaxPrefix(word)
if (maxPrefix.isEmpty) {
return None
}
val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
if (keyVal._1 != maxPrefix) {
return None
}
keyVal._2.get(word.substring(maxPrefix.length))
}
def contains(word: String): Boolean = {
return get(word) != None
}
def print(tabs: Int = 0) {
if (endPoint) println(" " * tabs + "<" + data.get + ">")
for (elem <- leaves) {
println(" " * tabs + elem._1)
elem._2.print(tabs + elem._1.length)
}
}
}
var t = new MinimizedTrie[Int](false, None)
t = t.add("aba", 0)
t = t.add("abab", 1)
t = t.add("a", 2)
t.print()
println()
t = t.remove("abab")
t.print()
println("aba: " + t.get("aba"))