In [1]:
class SuffixTrie(object):
    """ Encapsulates a suffix trie of a provided string t """
    
    def __init__(self, t):
        """ Make suffix trie from t """
        t += '$'  # terminator symbol
        self.root = {}
        for i in range(len(t)):  # for each suffix
            cur = self.root
            for c in t[i:]:  # for each character in i'th suffix
                if c not in cur:
                    cur[c] = {}  # add outgoing edge if necessary
                cur = cur[c]
    
    def follow_path(self, s):
        """ Follow path given by characters of s.  Return node at
            end of path, or None if we fall off. """
        cur = self.root
        for c in s:
            if c not in cur:
                return None  # no outgoing edge on next character
            cur = cur[c]  # descend one level
        return cur
    
    def has_substring(self, s):
        """ Return true if s appears as a substring of t """
        return self.follow_path(s) is not None
    
    def has_suffix(self, s):
        """ Return true if s is a suffix of t """
        node = self.follow_path(s)
        return node is not None and '$' in node
    
    def to_dot(self):
        """ Return dot representation of trie to make a picture """
        lines = []
        def _to_dot_helper(node, parid):
            childid = parid
            for c, child in node.items():
                lines.append('  %d -> %d [ label="%s" ];' % (parid, childid+1, c))
                childid = _to_dot_helper(child, childid+1)
            return childid
        lines.append('digraph "Suffix trie" {')
        lines.append('  node [shape=circle label=""];')
        _to_dot_helper(self.root, 0)
        lines.append('}')
        return '\n'.join(lines) + '\n'
In [2]:
strie = SuffixTrie('there would have been a time for such a word')
In [3]:
strie.has_substring('nope')
Out[3]:
False
In [4]:
strie.has_substring('would have been')
Out[4]:
True
In [5]:
strie.has_substring('such a word')
Out[5]:
True
In [6]:
strie.has_suffix('would have been')
Out[6]:
False
In [7]:
strie.has_suffix('such a word')
Out[7]:
True
In [8]:
# might have to do this first:
# %install_ext https://raw.github.com/cjdrake/ipython-magic/master/gvmagic.py
%load_ext gvmagic
In [9]:
%dotstr SuffixTrie("CAT").to_dot()
Suffix trie 0 1 0->1 A 4 0->4 C 8 0->8 T 10 0->10 $ 2 1->2 T 3 2->3 $ 5 4->5 A 6 5->6 T 7 6->7 $ 9 8->9 $