add go trie code

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