Python數據結構——簡單創建二叉樹並打印樹狀圖

爲了方便刷LeetCode中的二叉樹題，編寫了如下程序，用於創建二叉樹，並打印出樹狀圖，深度要求小於等於4.

程序如下：

# 創建二叉樹節點
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


# 給定二叉樹的前序遍歷和中序遍歷，獲得該二叉樹
def getBSTwithPreTin(pre, tin):
    if len(pre) == 0 | len(tin) == 0:
        return None
    root = TreeNode(pre[0])
    for order, item in enumerate(tin):
        if root .val == item:
            root.left = getBSTwithPreTin(pre[1:order+1], tin[:order])
            root.right = getBSTwithPreTin(
                pre[order+1:], tin[order+1:])
            return root


def listcreattree(root, llist, i):  # 用列表遞歸創建二叉樹，
    # 它其實創建過程也是從根開始a開始，創左子樹b，再創b的左子樹，如果b的左子樹爲空，返回none。
    # 再接着創建b的右子樹，
    if i < len(llist):
        if llist[i] == '#':
            return None  # 這裏的return很重要
        else:
            root = TreeNode(llist[i])
            # 往左遞推
            root.left = listcreattree(
                root.left, llist, 2*i+1)  # 從根開始一直到最左，直至爲空，
            # 往右回溯
            root.right = listcreattree(
                root.right, llist, 2*i+2)  # 再返回上一個根，回溯右，
            # 再返回根'
            return root  # 這裏的return很重要
    return root


# 二叉樹按層序轉換爲列表
def treeArray(pRoot):
    if not pRoot:
        return []
    resultList = []
    curLayer = [pRoot]
    while curLayer:
        curList = []
        nextLayer = []
        for node in curLayer:
            if node == '.':
                curList.append('.')
            else:
                curList.append(node.val)
                if node.left:
                    nextLayer.append(node.left)
                else:
                    nextLayer.append('.')
                if node.right:
                    nextLayer.append(node.right)
                else:
                    nextLayer.append('.')
        resultList.append(curList)
        curLayer = nextLayer
    return resultList


# 打印爲樹狀圖
def toTree4(t):
    n = len(t) - 1
    for i in range(1, n-1):
        for j in range(2*i):
            if t[i][j] == '.':
                t[i+1].insert(j+1, '.')
                t[i+1].insert(j+1, '.')
    result = []
    result.append('       {}       '.format(t[0][0]))
    result.append('    /     \\    ')
    result.append('   {}       {}   '.format(t[1][0], t[1][1]))
    result.append('  / \\     / \\  ')
    result.append(' {}   {}   {}   {} '.format(*t[2]))
    result.append('/ \\ / \\ / \\ / \\')
    result.append(' '.join([str(i) for i in t[3]]))
    for i in result[:2*n-1]:
        print(i)


# 深度小於等於3
def toTree3(t):
    n = len(t) - 1
    for i in range(1, n-1):
        for j in range(2*i):
            if t[i][j] == '.':
                t[i+1].insert(j+1, '.')
                t[i+1].insert(j+1, '.')
    result = []
    result.append('   {}   '.format(t[0][0]))
    result.append('  / \\  ')
    result.append(' {}   {} '.format(t[1][0], t[1][1]))
    result.append('/ \\ / \\')
    result.append('{} {} {} {}'.format(*t[2]))
    for i in result[:2*n-1]:
        print(i)


# 將上兩個函數合併
def printTree(tree):
    a = treeArray(tree)
    if len(a) < 5:
        toTree3(a)
    else:
        toTree4(a)

測試一下：

出自LeetCode第101題;

給定一個二叉樹，檢查它是否是鏡像對稱的。

例如，二叉樹 [1,2,2,3,4,4,3] 是對稱的。

1
   / \
  2   2
 / \ / \
3  4 4  3

 但是下面這個 [1,2,2,null,3,null,3] 則不是鏡像對稱的:

1
   / \
  2   2
   \   \
   3

import tree


def isSymmetric1(root):
    """
    :type root: TreeNode
    :rtype: bool
    """
    def helper(root, mirror):
        if not root and not mirror:
            return True
        if root and mirror and root.val == mirror.val:
            return helper(root.left, mirror.right) and helper(root.right, mirror.left)
        return False
    return helper(root, root)


def isSymmetric2(root):
    def isSame(p, q):
        if p and q and p.val == q.val:
            return isSame(p.left, q.right) and isSame(p.right, q.left)
        if not p:
            return not q
        return False
    if not root:
        return True
    return isSame(root.left, root.right)


a = [1, 2, 2, 3, 4, 4, 3]

t1 = tree.listcreattree(None, a, 0)
tree.printTree(t1)
print(isSymmetric2(t1))

b = [1, 2, 2, '#', 3, '#', 3]

t2 = tree.listcreattree(None, b, 0)
tree.printTree(t2)
print(isSymmetric2(t2))

結果如下：