数据结构、Python——实现一个Python风格的简单二叉树类，支持先中后层序遍历以及根据遍历重建二叉树方法

简介

简单的二叉数类，支持结构化字符串和中序遍历+任意一种其他遍历方式来建立一颗Python风格的二叉树。

所谓结构化字符串，即以ROOT[(LEFT[,RIGHT])]来表示一个节点和其左右子树结构，无右子树可省略逗号，无子树可省略括号，如0(1(,2),3(4,5(6(7,8),9)))中根节点为0，左子树为1(,2)，右子树为3(4,5(6(7,8),9))。结构如下：

所谓Python风格，实现了一些特殊方法，让二叉数可以被for等语句迭代。

测试代码

bt = BinTree('0(1(,2),3(4,5(6(7,8),9)))', data=range(10))
print(bt)
print(list(bt.preorder()))
print(list(bt.inorder()))
print(list(bt.postorder()))
print(list(bt))
rb_by_pre = BinTree.from_order(bt.inorder(), preorder=bt.preorder(), data=range(10))
rb_by_pos = BinTree.from_order(bt.inorder(), postorder=bt.postorder(), data=range(10))
rb_by_lev = BinTree.from_order(bt.inorder(), levelorder=bt.levelorder(), data=range(10))
print(rb_by_pre, rb_by_pos, rb_by_lev)
print(bt == rb_by_pre, bt == rb_by_pos, bt == rb_by_lev)
bt.show()

测试效果

代码

class Node:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


class BinTree:
    def __init__(self, treestr='', data=None):
    	""" 接受格式化字符串和可索引的data数据，如果传入data则格式化字符串中的节点名将作为索引到data中寻找对应值作为节点值，否则自身即为节点值 """
        if not self._treestr_vaild(treestr):
            raise TypeError('{!r} is not a vaild treestr'.format(treestr))
        self.root = self._build_tree(treestr, data)

    def __eq__(self, other):
        return bool(self.root) and bool(other.root) and all(x == y for x, y in zip(self, other))

    def __iter__(self):
        return self.levelorder()

    def __str__(self):
        return '{}({})'.format(self.__class__.__name__, self.get_treestr())

    def get_treestr(self):
        """ 返回二叉树格式化字符串 """
        return self._build_str(self.root).replace('(,)', '').replace(',)', ')')

    def show(self):
        """ 左转90°打印二叉树 """
        stack = []
        node = self.root
        lv = 0
        while node or stack:
            if node:
                stack.append((lv, node))
                lv += 1
                node = node.right
            else:
                lv, node = stack.pop()
                print(lv * '   ' + str(node.val))
                node = node.left
                lv += 1

    def levelorder(self):
        """ 按层遍历二叉树的迭代器 """
        queue = [self.root]
        while queue:
            node = queue.pop(0)
            if node:
                if node.left:
                    queue.append(node.left)
                if node.right:
                    queue.append(node.right)
                yield node.val

    def preorder(self):
        """ 先序遍历二叉树的迭代器 """
        stack = [self.root]
        while stack:
            node = stack.pop()
            if node:
                if node.right:
                    stack.append(node.right)
                if node.left:
                    stack.append(node.left)
                yield node.val

    def inorder(self):
        """ 中序遍历二叉树的迭代器 """
        stack = []
        node = self.root
        while node or stack:
            if node:
                stack.append(node)
                node = node.left
            else:
                node = stack.pop()
                yield node.val
                node = node.right

    def postorder(self):
        """ 后序遍历二叉树的迭代器 """
        stack1 = [self.root]
        stack2 = []
        while stack1:
            node = stack1.pop()
            if node:
                if node.left:
                    stack1.append(node.left)
                if node.right:
                    stack1.append(node.right)
                stack2.append(node)
        while stack2:
            yield stack2.pop().val

    @classmethod
    def from_order(cls, inorder, preorder=None, postorder=None, levelorder=None, data=None):
        """ 返回根据中序遍历 +（先序/后序/层序）遍历建立的二叉树，无法建立时返回空树 """
        try:
            return cls(cls._from_order(inorder, preorder, postorder, levelorder), data)
        except ValueError:
            return cls()

    @classmethod
    def _from_order(cls, inorder, preorder=None, postorder=None, levelorder=None):
        if preorder:
            inorder = list(inorder)
            preorder = list(preorder)
            root = preorder[0]
            root_index = inorder.index(root)
            left = cls._from_order(inorder[:root_index], preorder=preorder[1: root_index + 1])
            right = cls._from_order(inorder[root_index + 1:], preorder=preorder[root_index + 1:])
            return '{}({},{})'.format(root, left, right)
        elif postorder:
            inorder = list(inorder)
            postorder = list(postorder)
            root = postorder[-1]
            root_index = inorder.index(root)
            left = cls._from_order(inorder[:root_index], postorder=postorder[: root_index])
            right = cls._from_order(inorder[root_index + 1:], postorder=postorder[root_index: -1])
            return '{}({},{})'.format(root, left, right)
        elif levelorder:
            inorder = list(inorder)
            levelorder = list(levelorder)
            preorder = []

            def pre(low, high):
                index = 0
                for n in levelorder:
                    try:
                        index = inorder[low: high + 1].index(n) + low
                    except ValueError:
                        continue
                    preorder.append(n)
                    break
                if index > low:
                    pre(low, index - 1)
                if index < high:
                    pre(index + 1, high)
            pre(0, len(inorder) - 1)
            return cls._from_order(inorder, preorder=preorder)
        else:
            return ''

    def _build_tree(self, treestr, data):
        if not treestr:
            return None
        root, left, right = self.treestr_split(treestr)
        if data is not None:
            root = data[int(root)]
        node = Node(root)
        node.left = self._build_tree(left, data)
        node.right = self._build_tree(right, data)
        return node

    def _build_str(self, node):
        if node is None:
            return ''
        return '{}({},{})'.format(node.val, self._build_str(node.left), self._build_str(node.right))

    @staticmethod
    def _treestr_vaild(treestr: str) -> bool:
        count = 0
        for c in treestr:
            if c == '(':
                count += 1
            elif c == ')':
                count -= 1
            if count < 0:
                return False
        if count == 0:
            return True
        else:
            return False

    @staticmethod
    def treestr_split(treestr: str) -> tuple:
        """ str: 'ROOT[(LEFT[,RIGHT])]' -> tuple: (ROOT, LEFT, RIGHT) """
        root_index = treestr.find('(')
        if root_index == -1:
            return treestr, '', ''
        root = treestr[:root_index]
        rest = treestr[root_index + 1: -1]
        count = 0
        for index, c in enumerate(rest):
            if c == '(':
                count += 1
            elif c == ')':
                count -= 1
            elif c == ',' and count == 0:
                left = rest[:index]
                right = rest[index + 1:]
                break
        else:
            left = rest
            right = ''
        return root, left, right