suffixtree.py

gistfile1.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

__author__ = 'Encore Hu, <huyoo353@126.com>'

class Suffix(object):
    def __init__(self, start, end):
        self.start = start
        self.end   = end

    def __repr__(self):
        return 'Suffix(%d,%d)' % (self.start, self.end)

    def __str__(self):
        return self.__unicode__().encode('utf-8')

    def __unicode__(self):
        return 'Suffix(%d,%d)' % (self.start, self.end)

class Node(object):
    def __init__(self):
        self.suffix_link = -1

    def __repr__(self):
        return 'Node(%d)' % self.suffix_link

    def __str__(self):
        return self.__unicode__().encode('utf-8')

    def __unicode__(self):
        return 'Node(%d)' % self.suffix_link

class Edge(object):
    def __init__(self, start_node, end_node, start_char_index, end_char_index):
        self.start_node = start_node
        self.end_node   = end_node
        self.start_char_index = start_char_index
        self.end_char_index   = end_char_index
        

    def __repr__(self):
        return 'Edge(%d,%d,%d,%d)' % (self.start_node, self.end_node,self.start_char_index, self.end_char_index)

    def __str__(self):
        return self.__unicode__().encode('utf-8')

    def __unicode__(self):
        return 'Edge(%d,%d,%d,%d)' % (self.start_node, self.end_node,self.start_char_index, self.end_char_index)

    def get_label(self):
        return ''

class SuffixTree(object):
    """A suffix tree for string matching. 
    Uses Ukkonen's algorithm for init.
    """
    def __init__(self, s):
        ''' init with unicode s, from left to right, char by char.
        if s is encoded as other encoding, should be decoded to unicode first.'''
        self.s = s
        
        
        self.root = Node()
        self.nodes=[self.root] # log the nodes inserted
        self.edges={}          # log every edge between nodes
        self.length=len(self.s)
        print 'constructing with unicode string:',s
        
        self.curr_end_index = 1 # current end index, used in _add_prefix func.

        #active point means: from node active_node, choose edge active_edge(startswith a char,
        #this stored the char), next new char will be inserted at position active_length.
        #active_edge max=count(charset),english a-z, active_edge max=26 or 52(case sensitive)
        #eg (0,'a',1) from root, choose edge starts with 'a'(there must be only one edge starts with 'a')
        #at next 1 postion, the new char will insert into the edge.
        #if active_edge=None, insert a new edge into the root node.
        self.active_point = (0,None,0) #(active_node, active_edge, active_length)
        self.remainder = 1
        for i in xrange(self.length):
            self._add_prefix(i)

    def __repr__(self):
        return 'SuffixTree(%s)' % repr(self.s)

    def __str__(self):
        return self.__unicode__().encode('utf-8')

    def __unicode__(self):
        print self.nodes
        print self.edges
        for key, edge in self.edges.items():
            print key,edge,self.s[edge.start_char_index:edge.end_char_index]
        return 'SuffixTree(%s)' % repr(self.s)

    def _add_prefix(self, char_index):
        print '-'*80
        char = self.s[char_index]
        print 'processing',char
        suffix=(char_index, self.curr_end_index)
        print suffix

        # before every step, remainder =1
        #self.remainder = 1  

        found = -1

        #auto extend exists edges
        for key,edge in self.edges.items():            
            self.edges[key].end_char_index = self.curr_end_index

        # insert a new edge for new char
        for x in xrange(self.remainder,0,-1):
            print 'active_point =',self.active_point
            print 'remainder =',self.remainder
            #print x,
            print 'remaind %d substring' % self.remainder,self.s[char_index+1-x:char_index+1]
            remaind_string = self.s[char_index+1-x:char_index+1]
            for key in self.edges:
                if self.active_point[1] == None:
                    #if key[1]==char:
                    if self.s[self.edges[key].start_char_index: self.edges[key].start_char_index+x] == remaind_string:
                        print 'found an edge begin with %s' % remaind_string
                        found = key[0]
                        print 'active_edge:%s->%s' % (self.active_point[1],remaind_string[0])
                        self.active_point=(self.active_point[0],remaind_string[0],self.active_point[2]+1)
                        self.remainder += 1
                        break
                else:# not root node , donot modify active_edge, active_length.
                    if self.s[self.edges[key].start_char_index: self.edges[key].start_char_index+x] == remaind_string:
                        print 'found an EXSIT edge begin with %s' % remaind_string
                        found = key[0]
                        self.active_point=(self.active_point[0],self.active_point[1],self.active_point[2])
                        self.remainder += 1
                        break

        if found != -1:
            # found an exists edge begin with the current char.
            # remainder +=1
            #self.active_point=(self.active_point[0],self.active_point[1],self.active_point[2]+1)
            #self.remainder += 1

            print 'active_point =',self.active_point
            print 'remainder =',self.remainder
        else:
            # the current char not found in any edge, so, insert it into suffix
            # tree(add a new edge)
            self.nodes.append(Node())
            edge=Edge(0,len(self.nodes)-1, char_index, self.curr_end_index)
            self.edges[(char_index, char)]= edge
            print 'insert edge',(char_index, char),edge
            #self.active_point=(self.active_point[0],self.active_point[1],self.active_point[2]-1)
            #self.remainder -= 1

        #if self.active_point[0]==0:
        #    self.nodes.append(Node())
        #    edge=Edge(0,len(self.nodes)-1, char_index, self.curr_end_index)
        #    self.edges[(char_index, self.s[char_index])]= edge
        
        if len(self.nodes)==1:#only root node
            # insert a node, add an edge with the char as the edge label.
            self.nodes.append(Node())
            edge=Edge(0, 1, char_index, self.curr_end_index) # 0 means the root node
            self.edges[(char_index,self.s[char_index])] = edge
        #else:
        #    for edge in self.edges:
        #        if edge[1]==self.s[char_index]:
        #            print 'found an edge begin with %s' % self.s[char_index]

        self.curr_end_index += 1
        print '-'*80

    def _add_edge(self,start_node,edge):
        self.edges[edge.start_node]=edge

    def find(self,substring):
        '''find substring in the suffix tree\'s edges, ending at leaves(edges).'''
        return -1

    def find_maxlength_substring(self):
        return (0,-1)

    def count(self,substring):
        '''find substring in the suffix tree\'s edges, ending at branches(nodes).'''
        return 0

def test_help():
    help(Suffix)
    help(Node)
    help(Edge)
    help(SuffixTree)
    print dir(SuffixTree)

def test_func():
    text='abcad'
    stree=SuffixTree(text)
    print stree

def test_text(text):
    print '='*66
    stree=SuffixTree(text)
    print stree
    print '='*66

if __name__ == '__main__':
    #test_help()
    #test_func()
    #test_text('abc') #yes
    #test_text('abcd') #yes
    test_text('abca') #yes
    test_text('abcab') #no
    #test_text('abcad') #yes
    #test_text('abcabx') #no
    #test_text('abcabxa') #no
    #test_text('abcabxab') #no
    #test_text('abcabxabc') #no
    #test_text('abcabxabcd') #no
    

use the algorithm of http://stackoverflow.com/questions/9452701/ukkonens-suffix-tree-algorithm-in-plain-english
not finished yet.