前缀树-TrieTree

# 只处理了ascii字符

local Node = {}
Node.__cname = "util.TrieTree.Node"
Node.__index = Node

function Node.new(v)
    local obj = {}
    setmetatable(obj, Node)
    obj:ctor(v)
    return obj
end

function Node:ctor(v)
    self.value = v
    self.passCount = 0
    self.endCount = 0 ---isWord?
    self.children = nil
end

function Node:GetChild(key)
    if nil == self.children then
        return nil
    end
    return self.children[key]
end

function Node:SetChild(key, child)
    if nil == self.children then
        if nil == child then return end
        self.children = {}
    end
    self.children[key] = child
end

function Node:DiffPassCount(diff)
    self.passCount = self.passCount + diff
    return self.passCount
end

function Node:DiffEndCount(diff)
    self.endCount = self.endCount + diff
    return self.endCount
end


local TrieTree = {}
TrieTree.__cname = "util.TrieTree"
TrieTree.__index = TrieTree

function TrieTree.new()
    local obj = {}
    setmetatable(obj, TrieTree)
    obj:ctor()
    return obj
end

function TrieTree:ctor()
    self.count = 0
    self.root = Node.new()
end

function TrieTree:Clear()
    self.count = 0
end

function TrieTree:GetCount()
    return self.count
end

function TrieTree:Add(word)
    local cur = self.root
    cur:DiffPassCount(1)

    for i=1,#word do
        local c = string.sub(word, i, i)
        local child = cur:GetChild(c)
        if nil == child then
            child = Node.new(c)
            self.count = self.count + 1
            cur:SetChild(c, child)
        end
        cur = child
        cur:DiffPassCount(1)
    end
    cur:DiffEndCount(1)
end

function TrieTree:Contains(word)
    return self:WordCount(word) > 0
end

---统计单词出现次数
function TrieTree:WordCount(word)
    local ct = self.count
    if ct < 1 then return 0 end

    local cur = self.root
    for i=1,#word do
        local c = string.sub(word, i, i)
        local child = cur:GetChild(c)
        if nil == child then
            return 0
        end
        cur = child
    end
    return cur.endCount
end

function TrieTree:ContainsPrefix(prefix)
    return self:PrefixCount(prefix) > 0
end

function TrieTree:PrefixCount(prefix)
    local ct = self.count
    if ct < 1 then return 0 end

    local cur = self.root
    for i=1,#prefix do
        local c = string.sub(prefix, i, i)
        local child = cur:GetChild(c)
        if nil == child then
            return 0
        end
        cur = child
    end
    return cur.passCount
end

function TrieTree:Remove(word)
    local ct = self.count
    if ct < 1 then return false end

    local cur = self.root
    local validWord = false

    for i=1,#word do
        local c = string.sub(word, i, i)
        local child = cur:GetChild(c)
        if nil ~= child then
            validWord = true
            child:DiffPassCount(-1)
            if 0 == child.passCount then
                cur:SetChild(c, nil)
                self.count = self.count - 1
            end
            cur = child
        else
            break
        end
    end

    if validWord then
        self.root:DiffPassCount(-1)
        cur:DiffEndCount(-1)
    end
end

function TrieTree:__tostring()
    local ct = self.count
    if ct < 1 then return "" end

    local strTb = {}
    local level = 1
    local stack = {}
    local cur = self.root
    table.insert(stack, { cur, level })

    local printLevel = level
    while #stack > 0 do
        local top = stack[1]
        table.remove(stack, 1)
        local node = top[1]
        local level = top[2]
        if level ~= printLevel then
            printLevel = level
            table.insert(strTb, "\n")
        elseif printLevel > 1 then
            table.insert(strTb, ",")
        end
        table.insert(strTb, tostring(node.value))

        if node.children then
            for k, v in pairs(node.children) do
                table.insert(stack, { v, level+1 })
            end
        end
    end
    return table.concat(strTb)
end

return TrieTree

 

【参考】

数据结构之Trie字典树_端碗吹水的技术博客_51CTO博客

算法数据结构：Trie 树 - 掘金 (juejin.cn)

Trie树_咕噜咕噜崩的博客-CSDN博客_trie树

数据结构篇——字典树（trie树） (shuzhiduo.com)

【游戏开发实战】手把手教你在Unity中使用lua实现红点系统（前缀树 | 数据结构 | 设计模式 | 算法 | 含工程源码）_林新发的博客-CSDN博客_红点系统实现