diff --git a/go/trie.go b/go/trie.go new file mode 100644 index 0000000..96d9d1b --- /dev/null +++ b/go/trie.go @@ -0,0 +1,88 @@ +package main + +import ( + "fmt" + "bufio" + "os" + "strings" +) + +type trie struct { + nodes map[rune]trie + terminal bool +} + +func new_trie() trie { + return trie{ nodes: map[rune]trie {} } +} + +func (t trie)insert(word []rune) trie { + if len(word) == 0 { + t.terminal = true + return t + } + + trie, ok := t.nodes[word[0]] + + if ok { + t.nodes[word[0]] = trie.insert(word[1:]) + return t + } else { + trie := new_trie() + trie = trie.insert(word[1:]) + t.nodes[word[0]] = trie + } + return t +} + +func (t *trie)search(word []rune) bool { + if len(word) == 0 { + return t.terminal + } + trie, ok := t.nodes[word[0]] + return ok && trie.search(word[1:]) +} + +func (t *trie)count_substrings(word []rune) int { + return t._count_substrings(word, 0) +} + +func (t *trie)_count_substrings(word []rune, acc int) int { + if len(word) == 0 { + return acc + } + + next, ok := t.nodes[word[0]] + if ok { + // if we want a word to count as a prefix of itself, remove len clause + if next.terminal && len(word) > 1 { + acc += 1 + } + return next._count_substrings(word[1:], acc) + } else { + return acc + } +} + +func main() { + root := new_trie() + reader := bufio.NewReader(os.Stdin) + word_list := []string{} + for { + raw, _, _ := reader.ReadLine() + line := string(raw) + line = strings.Trim(line, " ") + if line == "" { + break + } + root = root.insert([]rune(line)) + word_list = append(word_list, line) + } + + sum := 0 + for _, word := range word_list { + sum += root.count_substrings([]rune(word)) + } + + fmt.Println("total score: ", sum) +}