diff --git a/src/.DS_Store b/src/.DS_Store new file mode 100644 index 0000000..cbc9853 Binary files /dev/null and b/src/.DS_Store differ diff --git a/src/suffix_tree/node.py b/src/suffix_tree/node.py index 5889b08..dc5497a 100644 --- a/src/suffix_tree/node.py +++ b/src/suffix_tree/node.py @@ -1,6 +1,6 @@ """Node classes for a Generalized Suffix Tree.""" -from typing import Optional, Tuple, List, Set, Dict +from typing import Optional, Tuple, List, Set, Dict, Union from .util import Path, Id, Symbol, Symbols, debug from . import lca_mixin, util @@ -204,8 +204,8 @@ def post_order(self, f) -> None: f(self) def find_path( - self, S: Symbols, start: int, end: int - ) -> Tuple["Node", int, Optional["Node"]]: + self, S: Symbols, start: int, end: int, return_nodedepth: bool = False + ) -> Union[Tuple["Node", int, Optional["Node"]], Tuple["Node", int, Optional["Node"], int]]: """Find a path starting from this node. The path is absolute. @@ -217,6 +217,7 @@ def find_path( node: Node = self matched_len = self.depth max_len = end - start + node_depth = 0 while matched_len < max_len: # find the edge to follow @@ -231,13 +232,20 @@ def find_path( matched_len += 1 if matched_len < child.depth: # the path ends between node and child + if return_nodedepth: + return node, matched_len, None, node_depth return node, matched_len, child # we reached another node, loop node = child + node_depth += 1 else: # no edge to follow + if return_nodedepth: + return node, matched_len, None, node_depth return node, matched_len, None # path exhausted + if return_nodedepth: + return node, matched_len, None, node_depth return node, matched_len, None def split_edge(self, new_len: int, child: Node) -> "Internal": diff --git a/src/suffix_tree/tree.py b/src/suffix_tree/tree.py index 87af9c9..99bfe51 100644 --- a/src/suffix_tree/tree.py +++ b/src/suffix_tree/tree.py @@ -116,15 +116,23 @@ def add( assert isinstance(builder, Builder) builder.build(self.root, id_, itertools.chain(S, [UniqueEndChar(id_)])) - def find_path(self, path: Path) -> Tuple[Node, int, Optional[Node]]: + def find_path( + self, path: Path, return_nodedepth: bool = False, + ) -> Union[Tuple[Node, int, Optional[Node]], Tuple[Node, int, Optional[Node], int]]: """Find a path in the tree. See: :py:func:`suffix_tree.node.Internal.find_path` """ + if return_nodedepth: + return self.root.find_path( + path.S, path.start, path.end, return_nodedepth=True + ) return self.root.find_path(path.S, path.start, path.end) - def find(self, S: Symbols) -> bool: + def find( + self, S: Symbols, return_nodedepth: bool = False + ) -> Union[bool, Tuple[bool, int]]: """Find a sequence in the tree. :param Sequence S: a sequence of symbols @@ -136,9 +144,18 @@ def find(self, S: Symbols) -> bool: True >>> tree.find("abc") False + >>> tree.find("abx", return_depth=True) + (True, 2) """ path = Path(S) + + if return_nodedepth: + dummy_node, matched_len, dummy_child, nodedepth = self.find_path( + path, return_nodedepth=return_nodedepth + ) + return matched_len == len(path), nodedepth + dummy_node, matched_len, dummy_child = self.find_path(path) return matched_len == len(path)