added trie

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
+}

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)
+		}
+	}
+}