Python 数据结构 tree 树

其他

如果想使用for ... in ...操作来遍子节点，我们可以实现items()方法：

class TreeNode(object):

    def items(self):
        return self.child.items()

如果想使用系统的in操作符，来判断是否存在名字为name的子节点，

class TreeNode(object):

    def __contains__(self, item):
        return item in self.child

如果想得到当前节点中子节点的个数，可以使用系统的len()函数。
我们所要做的就是重写__len__()

注意：如果重写__len__()的话，最好同时重写__bool__()。
因为python在做布尔判断时，如果没有找到__bool__()的话，会使用__len__()来替代。
这样就导致如果没有子节点，当前节点的布尔返回False

这里我们定义__bool__()永远返回True，这样我们可以通过布尔判断来判断一个节点是否存在。

class TreeNode(object):

    def __len__(self):
        """return number of children node"""
        return len(self.child)

    def __bool__(self, item):
        """always return True for exist node"""
        return True

如果想把树结构打印出来，可以创建一个dump()方法。

class TreeNode(object):

    def dump(self, indent=0):
        """dump tree to string"""
        tab = '    '*(indent-1) + ' |- ' if indent > 0 else ''
        print('%s%s' % (tab, self.name))
        for name, obj in self.items():
            obj.dump(indent+1)

如果想把树结构保存到文件里，稍候参考本人另一篇关于序列化的文章

源代码

类代码和测试代码如下（python2.7和python3）

#!/usr/bin/python

from __future__ import unicode_literals  # at top of module
from __future__ import division, print_function, with_statement



class TreeNode(object):
    """The basic node of tree structure"""

    def __init__(self, name, parent=None):
        super(TreeNode, self).__init__()
        self.name = name
        self.parent = parent
        self.child = {}

    def __repr__(self) :
        return 'TreeNode(%s)' % self.name


    def __contains__(self, item):
        return item in self.child


    def __len__(self):
        """return number of children node"""
        return len(self.child)

    def __bool__(self, item):
        """always return True for exist node"""
        return True

    @property
    def path(self):
        """return path string (from root to current node)"""
        if self.parent:
            return '%s %s' % (self.parent.path.strip(), self.name)
        else:
            return self.name

    def get_child(self, name, defval=None):
        """get a child node of current node"""
        return self.child.get(name, defval)

    def add_child(self, name, obj=None):
        """add a child node to current node"""
        if obj and not isinstance(obj, TreeNode):
            raise ValueError('TreeNode only add another TreeNode obj as child')
        if obj is None:
            obj = TreeNode(name)
        obj.parent = self
        self.child[name] = obj
        return obj

    def del_child(self, name):
        """remove a child node from current node"""
        if name in self.child:
            del self.child[name]

    def find_child(self, path, create=False):
        """find child node by path/name, return None if not found"""
        # convert path to a list if input is a string
        path = path if isinstance(path, list) else path.split()
        cur = self
        for sub in path:
            # search
            obj = cur.get_child(sub)
            if obj is None and create:
                # create new node if need
                obj = cur.add_child(sub)
            # check if search done
            if obj is None:
                break
            cur = obj
        return obj

    def items(self):
        return self.child.items()

    def dump(self, indent=0):
        """dump tree to string"""
        tab = '    '*(indent-1) + ' |- ' if indent > 0 else ''
        print('%s%s' % (tab, self.name))
        for name, obj in self.items():
            obj.dump(indent+1)


if __name__ == '__main__':
    # local test
    print('test add_child()')
    root = TreeNode('') # root name is ''
    a1 = root.add_child('a1')
    a1.add_child('b1')
    a1.add_child('b2')
    a2 = root.add_child('a2')
    b3 = a2.add_child('b3')
    b3.add_child('c1')
    root.dump()
    # (root)
    #  |- a1
    #      |- b1
    #      |- b2
    #  |- a2
    #      |- b3
    #          |- c1


    print('test items()')
    for name, obj in a1.items():
        print(name, obj)
    # b1 TreeNode(b1)
    # b2 TreeNode(b2)


    print('test operator "in"')
    print("b2 is a1's child = %s" % ('b2' in a1))
    # b2 is a1's child = True


    print('test del_child()')
    a1.del_child('b2')
    root.dump()
    print("b2 is a1's child = %s" % ('b2' in a1))
    # (root)
    #  |- a1
    #      |- b1
    #  |- a2
    #      |- b3
    #          |- c1
    # b2 is a1's child = False


    print('test find_child()')
    obj = root.find_child('a2 b3 c1')
    print(obj)
    # TreeNode(c1)

    print('test find_child() with create')
    obj = root.find_child('a1 b1 c2 b1 e1 f1', create=True)
    print(obj)
    root.dump()
    # TreeNode(f1)
    # (root)
    # |- a1
    #     |- b1
    #         |- c2
    #             |- b1
    #                 |- e1
    #                     |- f1
    # |- a2
    #     |- b3
    #         |- c1

    print('test attr path')
    print(obj.path)
    # a1 b1 c2 b1 e1 f1