Tries

// Trie stores words letter by letter, so words that share a beginning share a path.
type node struct {
	next map[rune]*node
	end  bool // a word finishes exactly here
}

type Trie struct{ root *node }

func NewTrie() *Trie { return &Trie{root: &node{next: map[rune]*node{}}} }

func (t *Trie) Insert(word string) {
	cur := t.root
	for _, r := range word {
		child, ok := cur.next[r]
		if !ok { // 1. no edge for this letter yet: create it
			child = &node{next: map[rune]*node{}}
			cur.next[r] = child
		}
		cur = child // 2. walk down one letter
	}
	cur.end = true // 3. mark the word's last node
}

// walk follows prefix from the root and returns the node it ends on, or nil.
func (t *Trie) walk(prefix string) *node {
	cur := t.root
	for _, r := range prefix {
		cur = cur.next[r]
		if cur == nil {
			return nil // 4. the path breaks: nothing stored starts with this prefix
		}
	}
	return cur
}

// Search reports whether the exact word was inserted.
func (t *Trie) Search(word string) bool {
	n := t.walk(word)
	return n != nil && n.end
}

// StartsWith reports whether any inserted word begins with prefix.
func (t *Trie) StartsWith(prefix string) bool { return t.walk(prefix) != nil }

// Autocomplete returns up to limit words that start with prefix, in alphabetical order.
// Walk to the prefix node, then collect words below it (depth-first, letters sorted).
func (t *Trie) Autocomplete(prefix string, limit int) []string {
	start := t.walk(prefix)
	if start == nil || limit <= 0 {
		return nil
	}
	var out []string
	var dfs func(n *node, path []rune)
	dfs = func(n *node, path []rune) {
		if len(out) >= limit {
			return
		}
		if n.end {
			out = append(out, string(path))
		}
		letters := make([]rune, 0, len(n.next))
		for r := range n.next {
			letters = append(letters, r)
		}
		sort.Slice(letters, func(i, j int) bool { return letters[i] < letters[j] })
		for _, r := range letters {
			dfs(n.next[r], append(path, r))
		}
	}
	dfs(start, []rune(prefix))
	return out
}

// ReplaceWords: replace each word with the shortest root that is a prefix of it.
func ReplaceWords(roots []string, sentence string) string {
	t := NewTrie()
	for _, r := range roots {
		t.Insert(r)
	}
	shortest := func(word string) string {
		cur := t.root
		for i, r := range word {
			cur = cur.next[r]
			if cur == nil {
				break
			}
			if cur.end {
				return word[:i+len(string(r))] // the first complete root on the path is the shortest
			}
		}
		return word
	}
	out, start := "", 0
	for i := 0; i <= len(sentence); i++ {
		if i == len(sentence) || sentence[i] == ' ' {
			out += shortest(sentence[start:i])
			if i < len(sentence) {
				out += " "
			}
			start = i + 1
		}
	}
	return out
}

// WordDictionary: Search treats '.' as "any single letter".
type WordDictionary struct{ t *Trie }

func NewWordDictionary() *WordDictionary { return &WordDictionary{t: NewTrie()} }

func (w *WordDictionary) AddWord(word string) { w.t.Insert(word) }

func (w *WordDictionary) Search(word string) bool {
	var match func(n *node, rest []rune) bool
	match = func(n *node, rest []rune) bool {
		if len(rest) == 0 {
			return n.end
		}
		if rest[0] == '.' { // wildcard: try every child
			for _, child := range n.next {
				if match(child, rest[1:]) {
					return true
				}
			}
			return false
		}
		child := n.next[rest[0]]
		return child != nil && match(child, rest[1:])
	}
	return match(w.t.root, []rune(word))
}