102 lines
2.2 KiB
Go
102 lines
2.2 KiB
Go
package trie
|
|
|
|
import (
|
|
"sort"
|
|
)
|
|
|
|
// Node represents a node in the trie
|
|
type Node struct {
|
|
children map[rune]*Node
|
|
isEnd bool
|
|
}
|
|
|
|
// New creates a new empty trie root node
|
|
func New() *Node {
|
|
return &Node{
|
|
children: make(map[rune]*Node),
|
|
isEnd: false,
|
|
}
|
|
}
|
|
|
|
// Add adds a word to the trie
|
|
func (n *Node) Add(word string) {
|
|
current := n
|
|
for _, char := range word {
|
|
if _, exists := current.children[char]; !exists {
|
|
current.children[char] = &Node{
|
|
children: make(map[rune]*Node),
|
|
isEnd: false,
|
|
}
|
|
}
|
|
current = current.children[char]
|
|
}
|
|
current.isEnd = true
|
|
}
|
|
|
|
// countWords recursively counts the number of words in the subtree
|
|
func (n *Node) countWords() int {
|
|
count := 0
|
|
if n.isEnd {
|
|
count++
|
|
}
|
|
for _, child := range n.children {
|
|
count += child.countWords()
|
|
}
|
|
return count
|
|
}
|
|
|
|
// Count returns the number of words that have the given prefix
|
|
func (n *Node) Count(prefix string) int {
|
|
current := n
|
|
for _, char := range prefix {
|
|
if _, exists := current.children[char]; !exists {
|
|
return 0
|
|
}
|
|
current = current.children[char]
|
|
}
|
|
return current.countWords()
|
|
}
|
|
|
|
// PrefixCount represents a prefix and its count
|
|
type PrefixCount struct {
|
|
Prefix string
|
|
Count int
|
|
}
|
|
|
|
// collectPrefixes recursively collects all prefixes and their counts
|
|
func collectPrefixes(n *Node, currentPrefix string, results *[]PrefixCount) {
|
|
count := n.countWords()
|
|
if count > 0 {
|
|
*results = append(*results, PrefixCount{
|
|
Prefix: currentPrefix,
|
|
Count: count,
|
|
})
|
|
}
|
|
// Sort children by rune to ensure consistent traversal order
|
|
chars := make([]rune, 0, len(n.children))
|
|
for char := range n.children {
|
|
chars = append(chars, char)
|
|
}
|
|
sort.Slice(chars, func(i, j int) bool {
|
|
return chars[i] < chars[j]
|
|
})
|
|
for _, char := range chars {
|
|
collectPrefixes(n.children[char], currentPrefix+string(char), results)
|
|
}
|
|
}
|
|
|
|
// TopPrefixes returns a sorted list of prefixes by their counts
|
|
func TopPrefixes(n *Node) []PrefixCount {
|
|
var results []PrefixCount
|
|
collectPrefixes(n, "", &results)
|
|
|
|
sort.Slice(results, func(i, j int) bool {
|
|
if results[i].Count == results[j].Count {
|
|
return results[i].Prefix < results[j].Prefix
|
|
}
|
|
return results[i].Count > results[j].Count
|
|
})
|
|
|
|
return results
|
|
}
|