把二叉树打印成多行

题目描述

从上到下按层打印二叉树，同一层结点从左至右输出。每一层输出一行。

思路

在宽度优先遍历的基础上，需要记录二叉树每行的结点个数，每一行打印成一个列表，最后的结果是一个二维列表。

宽度优先遍历过程采用队列完成。下面分析如何记录每行的个数：用变量toPrint记录当前层需要打印的结点个数,nextLevel记录下一层需要打印的个数，当toPrint变为0时，将nextLevel的值赋给toPrint, nextLevel清零重计。

# -*- coding:utf-8 -*-
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
class Solution:
    # 返回二维列表[[1,2],[4,5]]
    def Print(self, pRoot):
       # 用一个队列，与宽度优先遍历不同的是，这道题还需要记录每行的个数
        if not pRoot:
            return []
        toPrint = 1 #当前层待打印节点数
        nextLevel = 0
        res = [] #总结果
        rowRes = [] #每一层
        ls = [pRoot]
        while ls:
            node = ls.pop(0)
            rowRes.append(node.val)
            toPrint -= 1
            if node.left:
                ls.append(node.left)
                nextLevel += 1
            if node.right:
                ls.append(node.right)
                nextLevel += 1
            if toPrint == 0:
                toPrint = nextLevel
                nextLevel = 0
                res.append(rowRes)
                rowRes = []
        return res

类似的题目还有之字形打印二叉树，只需要再添加一个标记变量来控制打印顺序，如下。

class Solution:
    def Print(self, pRoot):
        # write code here
        if not pRoot:
            return []
        res = []
        resRow = []
        ls = []
        ls.append(pRoot)
        #=====初始化3个标记变量==============
        mark = 1 #1:从左往右， -1:从右往左
        toPrint = 1
        if pRoot.left and pRoot.right:
            nextLevel = 2
        elif pRoot.left or pRoot.right:
            nextLevel = 1
        else:
            nextLevel = 0
        #============主体循环============
        while ls:
            curNode = ls.pop(0)
            resRow.append(curNode.val)
            toPrint -= 1
            if toPrint == 0:
                if mark > 0:
                    res.append(resRow[:])
                else:
                    resRow.reverse()
                    res.append(resRow[:])
                mark = -mark
                toPrint = nextLevel
                nextLevel = 0
                resRow = []
            if curNode.left:
                ls.append(curNode.left)
                nextLevel += 1
            if curNode.right:
                ls.append(curNode.right)
                nextLevel += 1
        return res