In [1]:
import (
    "fmt"
    "strings"
)
In [2]:
type Node struct {
    child map[byte] int
}

type SuffixTrie struct {
    t string
    nodes []*Node
}
In [3]:
func NewNode() *Node {
    node := new (Node)
    node.child = make(map[byte] int)
    return node
}

func NewSuffixTrie(t string) *SuffixTrie {
    st := new(SuffixTrie)
    st.nodes = append(st.nodes, NewNode())
    for i := 0; i < len(t); i++ {
        cur := 0
        for _, c := range []byte(t[i:]) {
            if node, ok := st.nodes[cur].child[c] ; ok {
                cur = node
            } else {
                st.nodes[cur].child[c] = len(st.nodes)
                cur = len(st.nodes)
                st.nodes = append(st.nodes, NewNode())
            }
        }
        st.nodes[cur].child['$'] = len(st.nodes)
        st.nodes = append(st.nodes, NewNode())
    }
    return st
}
In [4]:
func (st *SuffixTrie) followPath(s string) int {
    // Follow path given by characters of s.  Return node at
    // end of path, or -1 if we fall off
    cur := 0
    for _, c := range []byte(s) {
        if node, ok := st.nodes[cur].child[c] ; ok {
            cur = node
        } else {
            return -1
        }
    }
    return cur
}
In [5]:
func (st *SuffixTrie) hasSubstring(s string) bool {
    // Return true of s appears as a substring of t
    return st.followPath(s) >= 0
}
In [6]:
func (st *SuffixTrie) hasSuffix(s string) bool {
    // Return true of s appears as a substring of t
    cur := st.followPath(s)
    _, ok := st.nodes[cur].child['$']
    return cur >= 0 && ok
}
In [7]:
func (st *SuffixTrie) toDot() string {
    // Return dot representation of trie to make a picture
    lines := []string{"igraph \"Suffix trie\" {",
                      "  node [shape=circle label=\"\"];"}
    todo := []int{0}
    for len(todo) > 0 {
        node := todo[0]
        todo = todo[1:]
        for c, child := range st.nodes[node].child {
            lines = append(lines, fmt.Sprintf("  %d -> %d [ label=\"%c\" ];", node, child, c))
            todo = append(todo, child)
        }
    }
    lines = append(lines, "}")
    return strings.Join(lines, "\n")
}
In [8]:
strie := NewSuffixTrie("there would have been a time for such a word")
In [9]:
strie.hasSubstring("nope")
Out[9]:
false
In [10]:
strie.hasSubstring("would have been")
Out[10]:
true
In [11]:
strie.hasSubstring("such a word")
Out[11]:
true
In [12]:
strie.hasSuffix("would have been")
Out[12]:
false
In [13]:
strie.hasSuffix("such a word")
Out[13]:
true
In [14]:
strie2 := NewSuffixTrie("CAT")
strie2.toDot()
Out[14]:
igraph "Suffix trie" {
  node [shape=circle label=""];
  0 -> 1 [ label="C" ];
  0 -> 5 [ label="A" ];
  0 -> 8 [ label="T" ];
  1 -> 2 [ label="A" ];
  5 -> 6 [ label="T" ];
  8 -> 9 [ label="$" ];
  2 -> 3 [ label="T" ];
  6 -> 7 [ label="$" ];
  3 -> 4 [ label="$" ];
}

Currently these Go notebooks can't render diagrams; see Python version for the actual tree diagram.