近段时间需要使用trie树来加速like操作，在网上找了一圈发现没有可用的scala实现的trie树于是自己改了一版。
该版本根据java版的trie树改编而来，java版地址:https://www.jianshu.com/p/b5b8f00d0e55
实现代码如下：
import scala.collection.mutable.{HashMap, ListBuffer}
class TrieTree {
private var root:TrieNode = null
def this(words: List[String]) {
this()
root = new TrieNode
try words.foreach(addWord)
catch {
case e: Exception =>
e.printStackTrace()
}
}
/**
* 向树中添加单词
* @param word
*/
def addWord(word: String): Unit = {
if (word == null || word.length == 0) return
// 每次调用addWord都重新拿到全局根节点对象
var current = root
for(i <- 0 until word.length){
val code = word.charAt(i)
current = current.add(code)
}
current.end = true
}
/**
* 判断文本中是否有树中可匹配的词
* @param text 要参与匹配的文本
* @return
*/
def isContains(text: String): Boolean = {
if (text == null || text.length == 0) return false
// 获得前缀树
var current = root
// 从词的首位开始遍历
var index = 0
while (index < text.length){
// 如果在当前层找到当前字母，继续往下一层找
if (current.child.getOrElse(text.charAt(index),null) != null)
current = current.child.getOrElse(text.charAt(index),null)
else
{ // 如果在当前这一层找不到字符子节点，直接切到新的root该子节点下重新找
// 如果root下也没有该字母，继续返回root给下一个字母调用防止空指针
current = if (root.child.getOrElse(text.charAt(index),null) == null) root
else root.child.getOrElse(text.charAt(index),null)
}
// 判断是否存在的依据： 当前查找返回的节点对象是否是end标志
if (current.end) return true
index += 1
}
false
}
/**
* 获得匹配到的第一个词
* @param text 要进行匹配的文本
* @return
*/
def getContainsItemOne(text: String): String = {
if (text == null || text.length == 0) return null
var current = root
var index = 0
var startIndex = 0
var endIndex = 0
while (index < text.length){
if (current.child.getOrElse(text.charAt(index),null) != null) current = current.child.getOrElse(text.charAt(index),null)
else { // startIndexstartIndex在else条件中更新
// 有两种情况，如果在根节点都找不到当前字则从index+1开始，如果根节点存在该字，从index开始
if (root.child.getOrElse(text.charAt(index),null) == null) {
current = root
startIndex = index + 1
}
else {
current = root.child.getOrElse(text.charAt(index),null)
startIndex = index
}
}
if (current.end) {
endIndex = index
return text.substring(startIndex, endIndex + 1)
}
index += 1
}
null
}
/**
* 获得可以配置到的所有词
* @param text 要匹配的文本
* @return
*/
def getContainsItemAll(text: String): List[String] = {
val res = ListBuffer[String]()
if (text == null || text.length == 0) return res.toList
var current = root
var index = 0
var startIndex = 0
var endIndex = 0
while (index < text.length) {
if (current.child.getOrElse(text.charAt(index),null) != null) current = current.child.getOrElse(text.charAt(index),null)
else { // startIndexstartIndex在else条件中更新
// 有两种情况，如果在根节点都找不到当前字则从index+1开始，如果根节点存在该字，从index开始
if (root.child.getOrElse(text.charAt(index),null) == null) {
current = root
startIndex = index + 1
}
else {
current = root.child.getOrElse(text.charAt(index),null)
startIndex = index
}
}
if (current.end) {
endIndex = index
res += text.substring(startIndex, endIndex + 1)
if (current.isLeaf) { // 重置为root
current = root
// 重置startIndex
startIndex = endIndex + 1
}
}
index += 1
}
res.toList
}
/**
* 一个节点对象
* value： 当前节点存储的字母
* child： 当前节点的子节点信息 字母 -> 节点对象
* end: 是否是整个词的结尾
*/
private class TrieNode{
var value = 0
var child: HashMap[Char, TrieNode] = null
var end = false
def add(newChar: Char):TrieNode = {
if (child == null) this.child = new HashMap[Char, TrieNode]
// 找到对应字符的字典树
var t = child.getOrElse(newChar,null)
// 在map中查找是否已经存在字母
if (t == null) { // 不存在则新建一个节点对象
t = new TrieNode
// 给节点对象命名为该字母
t.value = newChar
child.put(newChar, t)
}
// 返回下一个节点
t
}
//判断本节点是否为叶子节点
def isLeaf: Boolean = child == null
}
}
object TrieTree{
def apply(words: List[String]): TrieTree = {
new TrieTree(words)
}
def main(args: Array[String]): Unit = {
val trieTree = TrieTree(List("区块链", "智慧城市", "央行", "是吗", "区块"))
val res0 = trieTree.isContains("吃饱了")
val res1 = trieTree.isContains("吃饱了是吗")
val res2 = trieTree.getContainsItemOne("我是央行区块链之王")
val res3 = trieTree.getContainsItemAll("我是央行区块链之王")
println(res0) //false
println(res1) //true
println(res2) //央行
println(res3) //List(央行, 区块, 区块链)
}
}
