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