import scala.collection.immutable.HashMap
class CritBitTree[T](nodes: HashMap[String, CritBitTree[T]], isExternalNode: Boolean, data: Option[T]) {
def this(isExternalNode: Boolean, data: Option[T]) = this(HashMap[String, CritBitTree[T]](), isExternalNode, data)
def findNodeWithPrefix(prefix: String): (String, CritBitTree[T]) = {
val nodeWIthCommonPrefix = nodes.find(node => node._1.startsWith(prefix)).get
(nodeWIthCommonPrefix._1, nodeWIthCommonPrefix._2)
}
def contains(bitString: String): Boolean = {
have(bitString) != None
}
def have(bitString: String): Option[T] = {
if (bitString.length == 0) {
if (isExternalNode) {
return data
}
else
{
return None
}
}
val commonPrefix = findCommonPrefix(bitString)
if (commonPrefix.length == 0) {
return None
}
val (existPrefix, subtree) = findNodeWithPrefix(commonPrefix)
if (existPrefix != commonPrefix) {
return None
}
subtree.have(bitString.substring(commonPrefix.length))
}
def findIndex(index: Int, newString: String, existPrefix: String): Int = {
if (index < newString.length && index < existPrefix.length && newString(index) == existPrefix(index)) {
findIndex(index + 1, newString, existPrefix)
}
else
index
}
private def findCommonPrefix(bitString: String): String = {
for (existPrefixes <- nodes.keys) {
if (findIndex(0, bitString, existPrefixes) != 0) {
return bitString.substring(0, findIndex(0, bitString, existPrefixes))
}
}
bitString.substring(0, 0)
}
def add(bitString: String, value: T): CritBitTree[T] = {
if (bitString.length == 0)
return new CritBitTree[T](nodes, true, Option[T](value))
val commonPrefix = findCommonPrefix(bitString)
if (commonPrefix.length == 0) {
return new CritBitTree[T](nodes + (bitString -> new CritBitTree[T](true, Option[T](value))), isExternalNode, data)
}
val (existPrefix, subtree) = findNodeWithPrefix(commonPrefix)
if (existPrefix == commonPrefix) {
return new CritBitTree[T](nodes + (existPrefix -> subtree.add(bitString.substring(existPrefix.length), value)), isExternalNode, data)
}
else
{
val endPrefix = existPrefix.substring(commonPrefix.length)
val middleVertex = new CritBitTree[T](HashMap(endPrefix -> subtree), false, None)
val newNodes = (nodes - existPrefix) + (commonPrefix -> middleVertex.add(bitString.substring(commonPrefix.length), value))
return new CritBitTree[T](newNodes, isExternalNode, data)
}
}
def remove(bitString: String): CritBitTree[T] = {
if (bitString.length == 0) {
return new CritBitTree(nodes, false, None)
}
val commonPrefix = findCommonPrefix(bitString)
if (commonPrefix.length == 0) {
return this
}
val (existPrefix, subtree) = findNodeWithPrefix(commonPrefix)
if (existPrefix == commonPrefix) {
val childNode = existPrefix -> subtree.remove(bitString.substring(existPrefix.length))
if (bitString.length == commonPrefix.length) {
val newNodes = nodes - existPrefix
new CritBitTree[T](newNodes, isExternalNode, data)
}
else
{
val nextNode = childNode._2.getLastNode()
val newNodes = (nodes - existPrefix) + ((existPrefix + nextNode._1) -> nextNode._2)
new CritBitTree[T](newNodes, isExternalNode, data)
}
}
else
{
this
}
}
private def getLastNode() = nodes.last
def print(indent: Int = 0) {
if (isExternalNode)
println(" " * indent + "data: " + data.get)
for (node <- nodes) {
println(" " * indent + node._1)
node._2.print(indent + node._1.length)
}
}
}
object CritBitTree {
def apply[T](): CritBitTree[T] = {
new CritBitTree[T](false, None)
}
}