leetcode-go/leetcode/medium/648.replace-words.go

type trie struct {
	word bool
	next [26]*trie
}

func (t *trie) insert(s string) {
	for _, r := range s {
		i := r - 'a'
		if t.next[i] == nil {
			t.next[i] = &trie{}
		}
		t = t.next[i]
	}
	t.word = true
}

func (t *trie) search(s string) string {
	var sb strings.Builder
	for _, r := range s {
		t = t.next[r-'a']
		if t == nil {
			return ""
		}
		sb.WriteRune(r)
		if t.word {
			break
		}
	}
	return sb.String()
}

func replaceWords(dict []string, sentence string) string {
	var t trie
	for i := range dict {
		t.insert(dict[i])
	}
	words := strings.Split(sentence, " ")
	for i := range words {
		if root := t.search(words[i]); root != "" {
			words[i] = root
		}
	}
	return strings.Join(words, " ")
}