105 lines
2.4 KiB
Nim
105 lines
2.4 KiB
Nim
## Trie implementation
|
|
## 2023/02/08
|
|
import options
|
|
|
|
const
|
|
AlphabetSize = 26
|
|
type
|
|
NodeValue = char # array[4, uint8]
|
|
Trie* = object
|
|
root: TrieNode
|
|
size: int # number of words on the trie
|
|
TrieNode = ref object
|
|
nodes: seq[TrieNode]
|
|
value: NodeValue
|
|
isTerminal: bool # true if this node is the end of a word
|
|
|
|
|
|
proc add(self: TrieNode, value: NodeValue): TrieNode =
|
|
result = TrieNode(
|
|
nodes: newSeqOfCap[TrieNode](AlphabetSize),
|
|
value: value,
|
|
)
|
|
self.nodes.add(result)
|
|
|
|
proc getNode(self: TrieNode, value: NodeValue): Option[TrieNode] =
|
|
assert self != nil
|
|
for n in self.nodes:
|
|
if n.value == value:
|
|
return some(n)
|
|
return none(TrieNode)
|
|
|
|
proc followPath(self: Trie, str: string): Option[TrieNode] =
|
|
var node = self.root
|
|
for c in str:
|
|
# echo "path: ", c
|
|
let tmp = node.getNode(c)
|
|
if tmp.isNone:
|
|
return none(TrieNode)
|
|
node = tmp.get()
|
|
return some(node)
|
|
|
|
func len*(self: Trie): int = self.size
|
|
|
|
proc insert*(self: var Trie, word: string): var Trie {.discardable.} =
|
|
result = self
|
|
var node = self.root
|
|
for c in word:
|
|
let nextNode = node.getNode(c)
|
|
if nextNode.isNone:
|
|
node = node.add(c)
|
|
else:
|
|
node = nextNode.get()
|
|
node.isTerminal = true
|
|
inc self.size
|
|
|
|
proc delete*(self: var Trie, word: string): var Trie {.discardable.} =
|
|
result = self
|
|
let terminalNode = self.followPath(word)
|
|
if terminalNode.isNone:
|
|
return
|
|
terminalNode.get.isTerminal = false
|
|
dec self.size
|
|
|
|
proc prefixExists*(self: Trie, prefix: string): bool =
|
|
self.followPath(prefix).isSome
|
|
|
|
proc searchWord*(self: Trie, str: string): bool =
|
|
let nodeOpt = self.followPath(str)
|
|
if nodeOpt.isNone:
|
|
return false
|
|
return nodeOpt.get.isTerminal
|
|
|
|
proc getWordsByPrefix*(self: Trie, prefix: string): seq[TrieNode] =
|
|
let nodeOpt = self.followPath(prefix)
|
|
if nodeOpt.isNone:
|
|
return @[]
|
|
for child in nodeOpt.unsafeGet.nodes:
|
|
result.add(child)
|
|
|
|
when isMainModule:
|
|
var t = Trie(root: new(TrieNode))
|
|
let words = [
|
|
"gnome",
|
|
"gnostic",
|
|
"genetics",
|
|
"genome",
|
|
"geek",
|
|
"greek",
|
|
"general",
|
|
"generalist",
|
|
"generator",
|
|
"thomas",
|
|
"though",
|
|
"through"
|
|
]
|
|
|
|
for w in words:
|
|
t.insert(w)
|
|
assert t.searchWord(w)
|
|
|
|
for w in ["general", "gno", "th", "agact"]:
|
|
echo "'", w, "' ", "exists? ", t.prefixExists(w)
|
|
|
|
for n in t.getWordsByPrefix("gen"):
|
|
echo "child value: ", n.value
|