from typing import Tuple

# char: 单字符,形如'h','a','c','k'
# word: 词,形如"hack"


class TrieNode(object):
    """
    字符节点
    """
    def __init__(self, char: str):
        self.char = char  # 存储的字符
        self.children = []  # 该节点的子节点
        self.word_finished = False  # 是否是词尾
        self.counter = 1  # 出现在word中的次数


def add(root, word: str):
    """
    向字典树中添加词
    """
    node = root  # node的初始值为root节点*
    for char in word:  # 遍历整个word
        found_in_child = False  # 标记当前char是否在node的子结点列表中
        for child in node.children:
            if child.char == char:  # node节点的子节点中包含当前的char
                child.counter += 1  # 该字符在word中出现的次数加一
                node = child  # node向下移动一位,即为当前字符
                found_in_child = True  # 标记为找到
                break
        if not found_in_child:
            # 当前字符不在node的子节点列表中,则新建一个节点加入到node的子节点列表中
            new_node = TrieNode(char)
            node.children.append(new_node)
            node = new_node  # node移至当前节点
    # 当word中所有字符被遍历完后,node即为最后一个字符,将其标记为词尾
    node.word_finished = True


def find_profix(root, prefix: str) -> Tuple[bool, int]:
    """
    检查并返回:
    1. 该前缀是否出现在添加过的词中
    2. 如果是,那么有多少词组包含该前缀
    """
    node = root
    if not root.children:  # 如果当前节点无子节点则直接返回(False,0)
        return False, 0
    for char in prefix:  # 遍历prefix中的字符
        char_not_found = True  # 标记当前字符是否出现在node的子节点列表中,默认为True
        for child in node.children:
            if child.char == char:  # 如果当前字符包含在node的子结点中
                char_not_found = False
                node = child  # node向前移一位,即移至当前字符的位置
                break
        if char_not_found:  # 如果当前字符不在node的子结点中则直接返回(False,0)
            return False, 0
    # 执行至此步则表示prefix的所有字符都出现在字典树中,直接返回True和最后一个字符的counter
    return True, node.counter


def get_all_words(root: TrieNode):
    """
    获取Trie中所有的词组
    """
    all_words = []
    current_word = []
    # 递归方式
    def tmp(root: TrieNode):
        for node in root.children:  # 遍历当前节点的所有子节点
            current_word.append(node.char)  # 将当前子节点加入到current_word中
            if node.word_finished:  # 如果当前字符为词尾
                all_words.append(current_word.copy())  # 将当前词组加入到all_words中
            if node.children:
                tmp(node)  # 当前子节点有子节点, 继续递归遍历
            else:
                current_word.pop()  # 当前子节点无子节点即为词尾,出栈
        if current_word:
            current_word.pop()
        return

    tmp(root)  # 从根节点开始
    return [''.join(word) for word in all_words]


if __name__ == "__main__":
    # 构造根节点和字典树
    root = TrieNode('*')
    add(root, "hackathon")
    add(root, "hack")

    # 在字典树中查找
    print(find_profix(root, "hac"))
    print(find_profix(root, 'hack'))
    print(find_profix(root, 'hackathon'))
    print(find_profix(root, 'ha'))
    print(find_profix(root, 'hammer'))
    print(get_all_words(root))

# 输出结果为:
# (True, 2)
# (True, 2)
# (True, 1)
# (True, 2)
# (False, 0)
# ['hack', 'hackathon']