import ( "fmt" "strings" ) type Node struct { child map[byte] int } type SuffixTrie struct { t string nodes []*Node } 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 } 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 } func (st *SuffixTrie) hasSubstring(s string) bool { // Return true of s appears as a substring of t return st.followPath(s) >= 0 } 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 } 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") } strie := NewSuffixTrie("there would have been a time for such a word") strie.hasSubstring("nope") strie.hasSubstring("would have been") strie.hasSubstring("such a word") strie.hasSuffix("would have been") strie.hasSuffix("such a word") strie2 := NewSuffixTrie("CAT") strie2.toDot()