import sys class OrderedDict(dict): def __init__(self): super().__init__() self.first_letters = [None for _ in range(27)] def __setitem__(self, key, value): super().__setitem__(key, value) try: self.first_letters[self.rank(key)] = value except IndexError: raise IndexError(f"Could not add item of key '{key}' since it is out of range of the rank function.") def __delitem__(self, key): super().__delitem__(key) self.first_letters[self.rank(key)] = None def ordered_items(self): return filter(lambda x: x is not None, self.first_letters) @staticmethod def rank(char): # result = 26 if char == "$" else ord(char) - 97 result = 0 if char == "$" else ord(char) - 96 assert result in range(0, 27) return result class Node: global_end = 0 num_splits = 0 all_nodes = [] string = "" def __init__(self, start, end): self.root = False self.start = start self.end = end self.children = OrderedDict() self.id = len(self.all_nodes) self.suffix_index = self.id - self.num_splits - 1 self.all_nodes.append(self) self.parent = None self.link = None def __str__(self): link_str = "" if self.link is None else f" -> {self.link.id}" if not self.root: j, i = self.tuple() return f"[{self.id}, {self.tuple()}, {self.string[j:i + 1]}{link_str}]" return f"[{self.id} root{link_str}]" def __repr__(self): return f"[{self.id}]" def print_tree(self, spaces=1): print(f"{self}") for edge in self.children: print(f" " * spaces, end="") self.get_child(edge).print_tree(spaces=spaces + 1) def first_char(self): return self.string[self.start] def get_child(self, char): if char in self.children: return self.children[char] return None def add_child(self, child): child.parent = self self.children[child.first_char()] = child return child def remove_child(self, child): self.children.pop(child.first_char()) @property def end_index(self): return self.tuple()[1] def tuple(self): if self.root: raise Exception("Can't get substring of root.") if self.end == "#": return self.start, self.global_end return self.start, self.end @property def edge_length(self): if self.root: return 0 else: start, end = self.tuple() return end - start + 1 def detach(self): self.parent.remove_child(self) self.parent = None class Point: def __init__(self, node, edge="", length=0): assert isinstance(node, Node) self.node = node self.edge = edge self.length = length def __repr__(self): return f"(Node {self.node.id}'s edge:'{self.edge}', {self.length} along.)" def is_explicit(self): # a.k.a. is not on an edge return self.edge == "" def set_node(self, node): self.node = node self.edge = "" self.length = 0 if not self.is_explicit(): print("WARNING: Node.set_node", file=sys.stderr) @property def edge_node(self) -> Node: return self.node.get_child(self.edge) def index_here(self): if self.is_explicit(): return 0 if self.node.root else self.node.start return self.edge_node.start + self.length - 1 def char_here(self): return Node.string[self.index_here()] def create_root(): assert len(Node.all_nodes) == 0 root = Node(None, None) root.root = True root.link = root return root def split_edge(split_point: Point): assert not split_point.is_explicit() edge = split_point.edge_node original = edge.tuple() edge.detach() Node.num_splits += 1 mediator = Node(original[0], original[0] + split_point.length - 1) mediator.suffix_index = None edge.start = original[0] + split_point.length assert edge.start <= edge.end_index mediator.add_child(edge) split_point.node.add_child(mediator) return mediator def pos(n: int): return max(n, 0) def do_phase(root: Node, active: Point, i, last_j, remainder): root_point = Point(root) Node.global_end += 1 did_rule_three = False j = last_j + 1 node_just_created = None while not did_rule_three and j <= i + 1: curr_char = Node.string[i] match = char_is_after(active, curr_char) if match: # print(3) remainder += 1 if node_just_created is not None: node_just_created.link = active.node active = skip_count(1, active, i) did_rule_three = True else: # print(2) if not active.is_explicit(): mediator = split_edge(active) mediator.add_child(Node(i, "#")) if node_just_created is not None: node_just_created.link = mediator node_just_created = mediator active.length -= 1 if active.length == 0: active.set_node(active.node) else: active.node.add_child(Node(i, "#")) if node_just_created is not None and node_just_created.link is None: node_just_created.link = active.node remainder = pos(remainder - 1) active.set_node(active.node.link) if remainder > 0: active = skip_count(remainder, Point(root), i - remainder) last_j = j j += 1 # print(active) # root.print_tree() return active, remainder, last_j def char_is_after(point: Point, char): if point.is_explicit(): return char in point.node.children else: if point.length == point.edge_node.edge_length: return Node.string[point.edge_node.start] == char else: # If not at the end of an edge # return Node.string[point.index_here() + point.length] == char return Node.string[point.index_here() + 1] == char def skip_count(num_chars, start_point: Point, index): incoming_length = -1 existing_length = 0 head = start_point chars_left = num_chars char = "" if not head.is_explicit(): incoming_length = head.edge_node.edge_length - head.length if num_chars < incoming_length: head.length += num_chars return head head.set_node(head.edge_node) chars_left -= incoming_length index += incoming_length # Node.string[i] if head.node.root else Node.string[head.node.end_index + 1] # assert head.node.end_index + 1 + chars_left < len(Node.string) while chars_left > 0: # assert head.node.end_index + 1 + chars_left < len(Node.string) direction = Node.string[index] next_node = head.node.get_child(direction) if next_node is None: raise IndexError(f"Attempted to traverse char\n '{direction}' at point {head}. ({index=})") incoming_length = next_node.edge_length if chars_left < incoming_length: break chars_left -= incoming_length index += incoming_length head.set_node(next_node) # direction = Node.string[index] if chars_left > 0: # Landed on an edge head.edge = Node.string[index] head.length = chars_left return head def ukkonen(string): string += "$" Node.string = string Node.global_end = 0 Node.num_splits = 0 Node.all_nodes.clear() n = len(string) remainder = 0 last_j = 1 root = create_root() root.add_child(Node(0, "#")) active = Point(root) for i in range(1, n): active, remainder, last_j = do_phase(root, active, i, last_j, remainder) return root if __name__ == "__main__": # ukkonen("DEFDBEFFDDEFFFADEFFB") ukkonen("abacabad") print("done") # ukkonen("abcbcbc$")