python编写树寻找单词的最短前缀

"""
字符串的前缀是从给定字符串的开头开始的子字符串。例如，"carbon"的前缀是 "c", "ca", "car", "carb", "carbo", 和 "carbon"。 我们认为空字符串不是任何字符串的前缀，
每个非空字符串都是它自己的前缀。我们希望能用单词的前缀来缩写单词。例如 "carbohydrate" 通常被缩写为 "carb"。给出一个单词的集合，你需要找到能够唯一标识每一个单词的最短前缀。
在下面的示例输入中，“carbohydrate”可以缩写为“ carboh”，但不能缩写为“ carbo”（或更短的名称），因为列表中还有其他以“ carbo”开头的词。
完全匹配将覆盖前缀匹配。例如，前缀“ car”与给定的单词“ car”完全匹配。因此，可以唯一地理解“ car”是“ car”的缩写，而不是“ carriage”或以“ car”开头的列表中任何其他单词的缩写。
输入格式:
第一行        整数 n(集合中单词的数量，至少等于2)。
之后的n行  每行包含一个由小写字母组成的单词。
输出格式：
n行。每一行都包含来自输入对应行的单词，后跟一个空格，以及唯一（无歧义）标识该单词的最短前缀。

这个题目的难点就是最后一步，判断这个单词是否是路上的某一个节点。。。。
"""


class Node:
    def __init__(self, data=None):
        self.data = data            # 当前节点上面的数据
        self.next_dict = dict()     # 当前节点指向下级节点的字典

        self.data_list = []         # 用来保存每一级的数据
        self.short_names = dict()

        self.has_word = False       # 标记这个节点上是否有单词落地

    def add(self, string, i=0, node=None, word_s=None):
        if node is None:
            node = self

        if i >= len(string):
            return
        else:   # 非最短并且是已经存在的单词中的单词时候，将这个节点标记为分叉路口
            if string[0:i] in word_s:
                index = word_s.index(string[0:i])
                del word_s[index]
                # print(string[0:i], '落地，被标记的字符', node.data)
                node.has_word = True

        if string[i] not in node.next_dict:
            next_node = Node(string[i])
            self.add_data_to_list(string[i], i)
            node.next_dict[string[i]] = next_node
        else:
            next_node = node.next_dict[string[i]]

        i += 1
        self.add(string, i, next_node, word_s)

    def add_data_to_list(self, ch, i):
        if len(self.data_list) < i+1:
            self.data_list.append([])

        self.data_list[i].append(ch)

    def print_tree(self):   # 遍历打印每一层
        for i in range(len(self.data_list)):
            print('正在打印第{}层的数据'.format(i), self.data_list[i])

    # 站在每个数据节点上往下面处理
    def deal_node_tree(self):

        # 处理根节点下面的每一个节点
        for key in self.next_dict:
            self.deal_next_node(self.next_dict[key], '')
        return

    # 递归处理每一个节点
    def deal_next_node(self, node, string):
        # 递归的出口之一，当数据为None的时候
        if node.data is None:
            return

        string += node.data     # 拼接字符串

        # 从当前节点向下看有没有分岔路口,如果没有的话那就可以用当前字符串表示这个单词的最短缩写了,然后退出当前函数
        flag, string_new = self.is_else_road(node, node.next_dict, string)
        if flag is False:
            self.short_names[string_new] = string
            return

        # 如果不是的话，需要逐个递归处理所有的节点
        for key in node.next_dict:
            self.deal_next_node(node.next_dict[key], string)

    # 递归观察是否有分叉路口，以及单词落地点，如果有的话返回True,没有返回False
    def is_else_road(self, node, dict_data, string):
        if node.has_word:
            # print('此时碰到了一个落地点:', string)
            return True, string

        dict_len = len(dict_data)
        if dict_len == 0:   # 当没有下级节点说明到底了
            # print('末尾节点：', string)
            return False, string
        elif dict_len == 1:   # 只有一个节点的时候需要递归往下面
            for key in dict_data:
                only_node = dict_data[key]

                string += only_node.data
                return self.is_else_road(only_node, only_node.next_dict, string)   # 直接将处理结果返回出去
        else:
            return True, string     # 有多个路口的时候直接返回True


def get_data():
    n = int(input())
    data = [input().strip() for i in range(n)]
    return n, data


def main():
    n, data = get_data()

    root = Node()

    # 制造一颗字符树
    for x in data:
        root.add(x, word_s=data[0:])

    # root.print_tree()
    root.deal_node_tree()

    for x in data:
        if x in root.short_names:
            print('{} {}'.format(x, root.short_names[x]))
        else:
            print('{} {}'.format(x, x))

    return


if __name__ == '__main__':
    # str = 'aaa'
    # print(str[0:4])
    main()