From 3d71b654280cc57bf4e7fb1f455bd3159ebd6cc9 Mon Sep 17 00:00:00 2001 From: Colin Henry Date: Sat, 5 Jul 2025 11:55:28 -0700 Subject: [PATCH] added trie --- container/trie/trie.go | 101 ++++++++++++++++++++++ container/trie/trie_test.go | 163 ++++++++++++++++++++++++++++++++++++ 2 files changed, 264 insertions(+) create mode 100644 container/trie/trie.go create mode 100644 container/trie/trie_test.go diff --git a/container/trie/trie.go b/container/trie/trie.go new file mode 100644 index 0000000..4ec16bb --- /dev/null +++ b/container/trie/trie.go @@ -0,0 +1,101 @@ +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 +} diff --git a/container/trie/trie_test.go b/container/trie/trie_test.go new file mode 100644 index 0000000..aa06cd6 --- /dev/null +++ b/container/trie/trie_test.go @@ -0,0 +1,163 @@ +package trie + +import ( + "testing" +) + +func TestNew(t *testing.T) { + root := New() + if root == nil { + t.Error("New() returned nil") + } + if root.children == nil { + t.Error("New() root node has nil children map") + } + if root.isEnd { + t.Error("New() root node should not be marked as end") + } +} + +func TestAddAndCount(t *testing.T) { + tests := []struct { + name string + words []string + prefix string + expected int + }{ + { + name: "empty trie", + words: []string{}, + prefix: "test", + expected: 0, + }, + { + name: "single word", + words: []string{"hello"}, + prefix: "hel", + expected: 1, + }, + { + name: "multiple words same prefix", + words: []string{"hello", "help", "hell"}, + prefix: "hel", + expected: 3, + }, + { + name: "prefix not found", + words: []string{"hello", "help", "hell"}, + prefix: "xyz", + expected: 0, + }, + { + name: "empty prefix", + words: []string{"hello", "help", "hell"}, + prefix: "", + expected: 3, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + root := New() + for _, word := range tt.words { + root.Add(word) + } + if got := root.Count(tt.prefix); got != tt.expected { + t.Errorf("Count(%q) = %d, want %d", tt.prefix, got, tt.expected) + } + }) + } +} + +func TestTopPrefixes(t *testing.T) { + tests := []struct { + name string + words []string + expected []PrefixCount + }{ + { + name: "empty trie", + words: []string{}, + expected: []PrefixCount{}, + }, + { + name: "single word", + words: []string{"hello"}, + expected: []PrefixCount{ + {Prefix: "", Count: 1}, + {Prefix: "h", Count: 1}, + {Prefix: "he", Count: 1}, + {Prefix: "hel", Count: 1}, + {Prefix: "hell", Count: 1}, + {Prefix: "hello", Count: 1}, + }, + }, + { + name: "multiple words", + words: []string{"hello", "help", "hell", "helicopter"}, + expected: []PrefixCount{ + {Prefix: "", Count: 4}, + {Prefix: "h", Count: 4}, + {Prefix: "he", Count: 4}, + {Prefix: "hel", Count: 4}, + {Prefix: "hell", Count: 2}, + {Prefix: "hello", Count: 1}, + {Prefix: "help", Count: 1}, + {Prefix: "heli", Count: 1}, + {Prefix: "helic", Count: 1}, + {Prefix: "helico", Count: 1}, + {Prefix: "helicop", Count: 1}, + {Prefix: "helicopt", Count: 1}, + {Prefix: "helicopte", Count: 1}, + {Prefix: "helicopter", Count: 1}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + root := New() + for _, word := range tt.words { + root.Add(word) + } + got := TopPrefixes(root) + if len(got) != len(tt.expected) { + t.Errorf("TopPrefixes() returned %d results, want %d", len(got), len(tt.expected)) + return + } + for i, want := range tt.expected { + if i >= len(got) { + break + } + if got[i].Prefix != want.Prefix || got[i].Count != want.Count { + t.Errorf("TopPrefixes()[%d] = {Prefix: %q, Count: %d}, want {Prefix: %q, Count: %d}", + i, got[i].Prefix, got[i].Count, want.Prefix, want.Count) + } + } + }) + } +} + +func TestTopPrefixesSorting(t *testing.T) { + root := New() + words := []string{"a", "ab", "abc", "abcd", "abcde"} + for _, word := range words { + root.Add(word) + } + + prefixes := TopPrefixes(root) + + // Verify sorting by count (descending) + for i := 1; i < len(prefixes); i++ { + if prefixes[i].Count > prefixes[i-1].Count { + t.Errorf("Prefixes not sorted by count: %v", prefixes) + } + } + + // Verify alphabetical sorting for equal counts + for i := 1; i < len(prefixes); i++ { + if prefixes[i].Count == prefixes[i-1].Count && prefixes[i].Prefix < prefixes[i-1].Prefix { + t.Errorf("Prefixes with equal counts not sorted alphabetically: %v", prefixes) + } + } +}