| 10 | |
| 11 | |
| 12 | class SuffixTree: |
| 13 | def __init__(self, text: str) -> None: |
| 14 | """ |
| 15 | Initializes the suffix tree with the given text. |
| 16 | |
| 17 | Args: |
| 18 | text (str): The text for which the suffix tree is to be built. |
| 19 | """ |
| 20 | self.text: str = text |
| 21 | self.root: SuffixTreeNode = SuffixTreeNode() |
| 22 | self.build_suffix_tree() |
| 23 | |
| 24 | def build_suffix_tree(self) -> None: |
| 25 | """ |
| 26 | Builds the suffix tree for the given text by adding all suffixes. |
| 27 | """ |
| 28 | text = self.text |
| 29 | n = len(text) |
| 30 | for i in range(n): |
| 31 | suffix = text[i:] |
| 32 | self._add_suffix(suffix, i) |
| 33 | |
| 34 | def _add_suffix(self, suffix: str, index: int) -> None: |
| 35 | """ |
| 36 | Adds a suffix to the suffix tree. |
| 37 | |
| 38 | Args: |
| 39 | suffix (str): The suffix to add. |
| 40 | index (int): The starting index of the suffix in the original text. |
| 41 | """ |
| 42 | node = self.root |
| 43 | for char in suffix: |
| 44 | if char not in node.children: |
| 45 | node.children[char] = SuffixTreeNode() |
| 46 | node = node.children[char] |
| 47 | node.is_end_of_string = True |
| 48 | node.start = index |
| 49 | node.end = index + len(suffix) - 1 |
| 50 | |
| 51 | def search(self, pattern: str) -> bool: |
| 52 | """ |
| 53 | Searches for a pattern in the suffix tree. |
| 54 | |
| 55 | Args: |
| 56 | pattern (str): The pattern to search for. |
| 57 | |
| 58 | Returns: |
| 59 | bool: True if the pattern is found, False otherwise. |
| 60 | """ |
| 61 | node = self.root |
| 62 | for char in pattern: |
| 63 | if char not in node.children: |
| 64 | return False |
| 65 | node = node.children[char] |
| 66 | return True |