Python Longest Repeat

Question

I am trying to find the longest repeated string in text with python, both quickly and space efficiently. I created an implementation of a suffix tree in order to make the processing fast, but the code fails on large files due to memory problems.

I am purposely not using any advanced libraries for this particular code.

import string
import sys


class SuffixTree(object):

    class Node(object):
        """ The suffix tree is made of Node objects, which contain
            a position of where that substring it represents is
            in the larger text, a suffix which is all of the un-matched
            characters of a substring, and a "out" which contains links
            to the other nodes
        """
        def __init__(self, position, suffix):
            self.position = position
            self.suffix = suffix
            self.out = {}

    def __init__(self, text):
        """ the constructor for the SuffixTree takes in the text """

        self.text = text
        self.max_repeat = 2              # setting this to 2 initially because repeats of length 1 are not interesting
        self.repeats = {}                # dictionary to hold the repeats as we find them
        self.root = self.Node(None, '')  # initialize a root node
        L = len(self.text)               # calculate the length of the text

        """ Loop through the range of the text"""

        test = {}
        for i in xrange(L, 0, -1):
            try:
                self.branch(self.root, self.text[i-1:] + "$", i-1, 0)
            except MemoryError:
                print "Memory Error at the " + str(i) + "th iteration!"
                return

    def printRepeats(self):
        """ printRepeats simply finds the maximum repeat of the repeats we collected while building the tree
            and then prints out the information.  It should be noted that currently if there are more than one
            repeated sequence which are both of the same maximum length, this function is only going to print 
            out one of them 

            Also I am going in and changing the positions to be the biological positions by adding one """

        max_repeat = max(self.repeats.iterkeys(), key=len)
        print "Max repeat is", len(max_repeat), ":", max_repeat, "at locations:",

        for position in self.repeats[max_repeat]:
            print position+1,
        print


    def branch(self, currNode, substring, pos, repeat):
        """ branch function takes a particular substring, and starts at the current node (the root node)
            it checks to see if the current node contains a link to another node with the value of the
            substring's first letter.  If it does, it sets current node equal to THAT node and repeats
            the process recursively.  If a node has a suffix (leftovers that weren't needed to be matched)
            then it will create a new node based on the first letter in that suffix and attempt to see
            if it now has a branch it can follow.  When eventually it finds a location where it can get
            no deeper, it saves the current leftover letters as that node's suffix and returns recursively
            back to the main execution point (our constructor function).  

            Additionally while it is creating this branch, we are simultaneously creating the repeat
            information which we want to gather.  To do this, when it finds that it can no longer traverse
            down the node tree it will have maintained a count (variable called repeat) of how deep it has gone
            it will then save the position of where it stared as well as the positions of all of the adjacent
            level nodes into the repeat dictionary.  It will only do this however if the length of the repeat
            is longer than or equal to the maximum recorded repeat.  So for instance if we had found a repeat
            of length 3, it isn't going to bother saving any information about a repeat of position 2.
        """

        # check if the current node has leftover letters (called a suffix) if so, create a new branch
        if currNode.suffix != '':
            currNode.out[currNode.suffix[0]] = self.Node(currNode.position, currNode.suffix[1:])
            currNode.suffix = ''
            currNode.position = None


        # if there is no more paths that we can follow    
        while substring[repeat] in currNode.out:

            currNode = currNode.out[substring[repeat]]

            # check if the current node has leftover letters (called a suffix) if so, create a new branch
            if currNode.suffix != '':
                currNode.out[currNode.suffix[0]] = self.Node(currNode.position, currNode.suffix[1:])
                currNode.suffix = ''
                currNode.position = None
            repeat += 1
            #substring = substring[1:]


        else:

            # create a new node with its first letter, position, and put the rest of the letters in the suffix
            currNode.out[substring[repeat]] = self.Node(pos, buffer(substring,repeat+1))

            # check to see if the length of this repeat is >= the biggest ones we've found so far
            if repeat >= self.max_repeat:

                # go through each node at this branch and save its info to the repeat dictionary
                for node in currNode.out:
                    self.repeats.setdefault(self.text[pos:pos+repeat], []).append(currNode.out[node].position)

                # set the new maximum repeat size to this repeat size
                self.max_repeat = repeat




    text = readFile("ecoli.fasta")
    tree = SuffixTree(text)
    tree.printRepeats()

The file 'ecoli.fasta' is simply a text file that is 4.5MB large filled with A,G,T,C characters (it's a DNA sequence). The main for loop dies with a Memory Error after the 4,577,890th iteration. At this point there are approximately 63,762 Nodes.

Any suggestions?

Answer 1

You are using a buffer at one point:

        currNode.out[substring[repeat]] = self.Node(pos, buffer(substring,repeat+1))

That's good because it avoids copying the string. However, that does not save you because substring originates from the expression self.text[i-1:] + "$" which does create a copy. Likewise, slicing the buffer further in currNode.suffix[1:] creates a str object that copies the buffer.

It would seem to me that you could simply change these to use buffers as well (though I did not test; some other adjustments may be necessary):

    text_dollar = self.text + "$"
    for i in xrange(L, 0, -1):
        try:
            self.branch(self.root, buffer(text_dollar, i-1), i-1, 0)

and

    if currNode.suffix:
        currNode.out[currNode.suffix[0]] = self.Node(currNode.position, buffer(currNode.suffix, 1))
        currNode.suffix = ''
        currNode.position = None