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="$" ];
}