2
\$\begingroup\$

Here is my attempt to implement the Ukkonen algorithm for the Finding a Shared Motif problem on rosalind.info.

The code works, it produces a correct answer. It does so in 11 seconds. I am hoping to optimize it so that it runs faster.

# http://rosalind.info/problems/lcsm/

from sys import stdin

# A suffix tree is a tree that represents all the suffixes of a string.
# This implementation allows the string to be extended at the end.

# Conceptually, the edges between the nodes has some characters on it.
# If one reads from the root node to a leaf node, one should read a
# suffix of the string. If the last character is unique, the tree
# contains all the suffixes.

# In any tree, except the root, every node has a parent edge, therefore
# We just store the edge information with the root. Because we have the
# string, and edge labels are always substrings, it is enough to just
# store the indices.

# The parent and suffix link fields are a bit unfortunate. They are 
# artifacts needed for the construction of the tree, but otherwise
# often unnecessary for using the tree.

class suffix_tree_node:

    def __init__(self):
        self._begin = 0
        self._end = 0
        self._parent = None
        self._first_child = None
        self._sibling = None
        self._suffix_link = None

    def get_child(self):
        return self._first_child

    def get_sibling(self):
        return self._sibling

#
# Here is an implementation of the Ukkonen's suffix tree construction algorithm
# The code is ported from my C++ implementation.
#
# The code is best understood by reading Dan Gusfield's book, because that's what
# I read and implemented it based on the book. It is chapter 6.1
#
# At a high level, the Ukkonen's algorithm allows growing the suffix tree one 
# character at a time. When the jth character is added, conceptually add all the 
# new suffixes to the tree. 
#
# Now start reading the append() method to understand how it is achieved.
#
class suffix_tree:

    def __init__(self):
        self._root = suffix_tree_node()
        self._last_internal_node = None
        self._start = 0
        self._s = []

    #
    # To add a character to a suffix tree, we add all the new suffixes to the tree
    # A suffix starting from _start is added to the tree by calling the extension method
    # 
    # To be consistent with the book, we will call each invocation of this method a phase.
    # 
    # Conceptually, _start should start from 1. It is not in the code because of the 
    # following rule:
    #
    # Once a leaf, always a leaf.
    # 
    # Suppose the last phase introduced a leaf node, in this phase, the only thing that we
    # need to do for this phase is to extend the edge label. All leaves must terminate at the 
    # end of the string, so if we implicity interprets the end of a leaf node to be the length
    # of the string and never explicitly store it, there is nothing we need to for them. Therefore
    # we skipped them all.
    #
    # Here begs the question, how do we know from 0 to _start - 1 were leaves in the last phase,
    # we will read more about that in the extension() method.
    #
    # Another things to notice is that if the already_in_tree returned by the extension method is True,
    # then we stop adding. There is because of the following rule:
    # 
    # Rule 3 is a show stopper
    #
    # Without referring to the book, this would be a mystery. What is rule 3 anyway? Here is an 
    # attempt to describe it succinctly. In the following, We use capital letter to represent strings
    # and lowercase character to represent characters. 
    #
    # Suppose we know that in the last extension, we were trying to insert aXb into the tree and we 
    # discovered that it is already there. That means aXb was a suffix in the last phase, so Xb must
    # also be a suffix in the last phase, so we can skip adding Xb. But not just that, any suffixes 
    # of Xb
    #
    # Another interesting piece in this code is the next_node_cursor and next_text_cursor. They can
    # be understood as the cursors where we search in the tree. We will soon see how they are used
    # in the extension method. For now, what we needed to know is that the extension() method have 
    # some way to speed up the next suffix insertion by maintaining these states.
    #
    # Now start reading the extension() method to see how a single suffix is added the tree.
    #
    def append(self, c):
        self._s.append(c)
        next_node_cursor = self._root
        next_text_cursor = self._start
        while (True):
            (already_in_tree, next_node_cursor, next_text_cursor) = self._extension(next_node_cursor, next_text_cursor)
            if already_in_tree:
                break
            else:
                self._start = self._start + 1
                if self._start == len(self._s):
                    break

    def get_root(self):
        return self._root

    def get_edge(self, node):
        length = self._length(node)
        return (node._begin, node._begin + length)

    #
    # The extension method add a suffix into the tree. We knew that the whole suffix besides the last character 
    # must be already in the tree during the last phase, therefore the first thing to do is to find out where
    # to add the last character. This is done by the search() method.
    #
    # Without looking into the search() method yet, it is useful to describe how do we represent a location in
    # the tree. One could imagine the tree labels are highlighted and the highlight stop somewhere. There are two
    # cases, either the highlighting is filling every edge label, or the highlighting fill the last edge partially.
    # In both case, we can represent the location by a pair, the node that owns the edge label, and the number of 
    # character filled by the edge label. That is how node_cursor and edge_cursor are interpreted.
    # 
    # text_cursor is simply the index of the next character to be added in the tree, so there is no need for a variable
    # for it, it is simply len(self._s).
    # 
    # the next_* variables are optimization. They allows the search to be short circuited so that it doesn't always 
    # start from the beginning. For now, just assume the search method always returns node_cursor and edge_cursor at
    # the right place.
    # 
    # Now here we are a few cases. The first branch is the case where the edge label is completely filled. If the node
    # is a leaf, there is nothing to do because the implicit interpretation of the leaf node will make sure the last 
    # character is added. That said, the already_in_tree flag stays False, because conceptually we still added the 
    # character to the tree.
    #
    # In case the edge label is completely filled and yet it is not a leaf, there are still two cases. It could be the 
    # case that there is a child node that continues with the last character. In that case we know the suffix is already
    # in the tree, so we set already_in_tree to True and stop. Otherwise we create a new leaf node and stop.
    #
    # Similarly, in case the edge label is not completely filled, we check if the last character is already in the tree.
    # If it does, we set already_in_tree to True and stop. Otherwise, we have to split the edge.
    #
    # In all cases, we maintain two variables. search_end and new_internal_node. search_end represents the last node we 
    # reached in the search, and new internal node represents the new internal node we created in the split edge case.
    # Theses two variables are used to build suffix links, a tool for speeding up searches.
    # 
    # Suffix links only applies to nodes that are not leaves and not root. If such a node represents the prefix 'aX', 
    # then the suffix link of it points to an internal node that represents X. Suppose during the insertion of 'aXY' 
    # found such a link, there we could use that to speed to the search for 'XY' in the next extension. It is obvious 
    # why it could be useful. What is not so obvious is that suffix links are easy to build and always available.
    #
    # A theorem in the book shows that if an internal node is created in the current extension, the next extension
    # would have found its suffix link. Therefore, we save the new_internal_node and set its suffix link to search_end
    # in the next iteration.
    #
    # To see how the suffix link is used to speed up the search, read the search() method now.
    #
    def _extension(self, next_node_cursor, next_text_cursor):
        node_cursor = next_node_cursor
        edge_cursor = self._length(node_cursor)
        already_in_tree = False
        (next_node_cursor, next_text_cursor, node_cursor, edge_cursor) = self._search(next_text_cursor, node_cursor, edge_cursor)
        next_text_char = self._s[len(self._s) - 1]
        search_end = None
        new_internal_node = None
        if edge_cursor == self._length(node_cursor):
            if (not (node_cursor == self._root)) and (node_cursor._first_child == None):
                pass
            else:
                search = node_cursor._first_child
                found = False
                while search != None:
                    if self._first_char(search) == next_text_char:
                        found = True
                        break
                    else:
                        search = search._sibling
                if found:
                    already_in_tree = True
                else:
                    new_leaf = suffix_tree_node()
                    new_leaf._begin = len(self._s) - 1
                    new_leaf._parent = node_cursor
                    new_leaf._sibling = node_cursor._first_child
                    node_cursor._first_child = new_leaf
            search_end = node_cursor
        else:
            next_tree_char = self._s[node_cursor._begin + edge_cursor]
            if next_text_char == next_tree_char:
                already_in_tree = True
            else:
                new_node = suffix_tree_node()
                new_leaf = suffix_tree_node()
                new_leaf._begin = len(self._s) - 1
                new_node._begin = node_cursor._begin
                new_node._end = node_cursor._begin + edge_cursor
                node_cursor._begin = node_cursor._begin + edge_cursor

                new_node._parent = node_cursor._parent
                new_leaf._parent = new_node
                node_cursor._parent = new_node

                new_node._sibling = new_node._parent._first_child
                new_node._parent._first_child = new_node

                search = new_node
                while not (search == None):
                    if (search._sibling == node_cursor):
                        search._sibling = search._sibling._sibling
                        break
                    search = search._sibling

                new_node._first_child = new_leaf
                new_leaf._sibling = node_cursor
                node_cursor._sibling = None

                new_internal_node = search_end = new_node

        if not (self._last_internal_node == None):
            self._last_internal_node._suffix_link = search_end
            self._last_internal_node = None

        if not (new_internal_node == None):
            self._last_internal_node = new_internal_node

        return (already_in_tree, next_node_cursor, next_text_cursor)

    #
    # The search method does two things, it starts from the node_cursor and edge_cursor
    # and find all the way the all but last character of the string. In the process, we
    # also prepare the next_node_cursor and next_text_cursor for the next search.
    #
    # The search for the location is pretty straightforward, there are two things worth
    # notice. When we hit an edge, we just simply moved the cursor without checking the 
    # character. It is because we knew it has to be correct. 
    # 
    # In the book, this implementation optimization is called the skip/count trick.
    #
    # We also set the next_node_cursor and next_text_cursor whenever we found a suffix link.
    # This is where the suffix link serve its purpose, it make the next search much faster 
    # by starting closest to where we already knew it must reach.
    #
    def _search(self, text_cursor, node_cursor, edge_cursor):
        next_node_cursor = self._root
        next_text_cursor = text_cursor + 1
        while (text_cursor < len(self._s) - 1):
            node_length = self._length(node_cursor)
            if (edge_cursor == node_length):
                if not (node_cursor._suffix_link == None):
                    next_node_cursor = node_cursor._suffix_link
                    next_text_cursor = text_cursor
                next_char = self._s[text_cursor]
                child_cursor = node_cursor._first_child
                while True:
                    if (self._first_char(child_cursor) == next_char):
                        node_cursor = child_cursor
                        edge_cursor = 0
                        break
                    else:
                        child_cursor = child_cursor._sibling
            else:
                text_move = len(self._s) - 1 - text_cursor
                edge_move = node_length - edge_cursor
                if text_move > edge_move:
                    move = edge_move
                else:
                    move = text_move
                edge_cursor = edge_cursor + move
                text_cursor = text_cursor + move
        return (next_node_cursor, next_text_cursor, node_cursor, edge_cursor)

    def _length(self, node):
        if node == self._root:
            return 0
        elif node._first_child == None:
            return len(self._s) - node._begin
        else:
            return node._end - node._begin

    def _first_char(self, node):
        return self._s[node._begin]

#
# This is a helper function for the longest_common_substring to perform a depth first search
# For each node, we can imagine that it represents the set of suffixes that is descendant 
# leaf nodes does.
#
# We would like to know, for each node, the set of suffixes contains all suffixes that 
# start with each DNA. The prefix represented by the node is then a candidate for the longest
# common substring.
#
# To compute that set of DNAs, we conceptually requires nodes to return what DNA do they have
# but if we do so, all nodes must combine all the answer. And the combination time would take
# time proportional to the number of children. That's not good. Instead, we can let the children
# fill in the blanks. If we hand the same blank paper to all children and each of them mark on 
# it, then the combination is done implicitly. The time required to spend on the nodes is not 
# dependent of the number of children, but just on the number of DNAs, that makes the algorithm
# O(NK), where N is the total length of the combined DNAs and K is the number of DNAs.
#
def longest_common_substring_find(tree, node, length, index, blanks, answer):
    (begin, end) = tree.get_edge(node)
    child = node.get_child()
    if child == None:
        begin = begin - length
        found = 0
        for i in range(0, len(index)):
            if begin <= index[i]:
                found = i
                break
        blanks[found] = True
    else:
        my_blanks = [False] * len(blanks)
        while not (child == None):
            longest_common_substring_find(tree, child, length + end - begin, index, my_blanks, answer)
            child = child.get_sibling()
        good = True
        for i in range(0, len(blanks)):
            blanks[i] = my_blanks[i] or blanks[i]
            good = good and my_blanks[i]
        if good:
            begin = begin - length
            length = end - begin
            if answer["max"] < length:
                answer["max"] = length
                answer["begin"] = begin
                answer["end"] = end

#
# To find the longest common substring of a collection of texts
# We concatenate the text together, using lowercase letters (which we know cannot be DNA characters)
# to separate them and build a suffix tree. The longest common substring must be a common prefix
# of all suffixes, so we will use the longest_common_substring_find to perform a depth first search
# for the answer
#
def longest_common_substring(texts):
    separator = 'a'
    count = 0
    tree = suffix_tree()
    index = []
    all = ""
    for text in texts:
        for c in text:
            tree.append(c)
            count = count + 1
        all = all + text + separator
        index.append(count)
        tree.append(separator)
        count = count + 1
        separator = chr(ord(separator) + 1)
    answer = {}
    answer["max"] = -1
    longest_common_substring_find(tree, tree.get_root(), 0, index, [False] * len(texts), answer)
    print(all[answer["begin"]:answer["end"]])

data = []
for line in stdin:
data.append(line.strip())
data.append(">")

records = []
label = None
dna = ""
for line in data:
if line[0] == '>':
    if label != None:
    records.append(dna)
    label = line[1:] 
    dna = ""
else:
    dna = dna + line

longest_common_substring(records)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.