merge_trees.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-


def pretty_print_tree(tree, indent=0):
    """print pretty tree str"""
    print(' '*indent + 'Tree(%s)' % (tree.key))
    if tree.children:
        print(' '*indent + 'children:')
        for child in tree.children:
            pretty_print_tree(child, indent+1)



class Tree(object):

    def __init__(self, key, data, relation=None, children=[]):
        """:param children: ('father', a), ('teacher', b)"""
        self.key = key
        self.data = data
        self.children = children
        self.relation = relation or None

    def __str__(self):
        ret = 'Tree(%s): \n' % (self.key)

        return ret

    def __repr__(self):
        ret = 'Tree(%s): Children:\t%s' % (self.key, self.children)

        return ret

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            if other.data == self.data:
                if (len(other.children) == len(self.children)):
                    for i, v in enumerate(self.children):
                        if not v == other.children[i]:
                            return False
                    return True
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __gt__(self, other):
        return self.key > other.key

    def __gte__(self, other):
        return not self.key < other.key

    def __lt__(self, other):
        return self.key < other.key

    def __lte__(self, other):
        return not self.key > other.key

    @property
    def is_leaf(self):
        return not self.children

    @property
    def is_one_level_tree(self):
        if not self.children:
            return False
        for child in self.children:
            if not child.is_leaf:
                return False
        return True


def break_tree(tree):
    broken_trees = []
    if tree.is_leaf:
        return []
    for child in tree.children:
        broken_trees.append(Tree(tree.key, tree.data, None, [Tree(child.key, child.data, child.relation)]))
        if child.is_leaf:
            continue
        else:
            broken_trees.extend(break_tree(child))
    return broken_trees


def break_trees(trees):
    """break trees into one level trees"""
    broken_trees = []
    for tree in trees:
        broken_tree = break_tree(tree)
        broken_trees.extend(broken_tree)

    return broken_trees


def main():
    """docstring for main"""
    # build test data
    trees = [
        Tree('a', '赵春来', None, [
            Tree('a1', '赵玉龙','父亲', [
                Tree('a11', '程度', '主人', [
                    Tree('a111', '程虎', '哥哥')
                ])
            ]),
            Tree('a2', '高玉良', '领导')
        ]),
        Tree('b', '梁书记', None, [
            Tree('b1', '梁璐', '父亲'),
            Tree('b2', '祁同伟', '岳父'),
        ]),
        Tree('c', '高玉良',  None, [
            Tree('c1', '陈海', '老师'),
            Tree('c2', '祁同伟', '老师'),
            Tree('c3', '侯亮平', '老师'),
        ]),
        Tree('d', '陈岩石', None, [
            Tree('d1', '陈海', '父亲'),
            Tree('d2', '沙瑞金', '叔叔')
        ]),
        Tree('e', '沙班长', None, [
            Tree('e1', '陈岩石', '入党介绍人', [
                Tree('e11', '陈海', '父亲')
            ]),
            Tree('e2', '沙瑞金', '叔叔'),
        ])
    ]

    for tree in trees:
        pretty_print_tree(tree)
    print('-'*20)
    broken_trees = sorted(break_trees(trees))
    # broken_trees = break_trees(trees)
    for broken_tree in broken_trees:
        pretty_print_tree(broken_tree)


if __name__ == '__main__':
    main()