跳转至

算法:二叉树基于栈遍历

image-20191203131742924

1、前序遍历

算法:

  1. 若栈非空输出栈顶节点,并出栈
  2. 将右节点压栈(如果存在)
  3. 将左节点压栈(如果存在)
  4. 重复第1步直到栈空

代码:

from collections import deque


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


def preorder_stack(root: TreeNode):
    if not root:
        return
    stack = deque()
    stack.append(root)

    while stack:
        current = stack.pop()
        # 操作
        print(current.val)
        if current.right:
            stack.append(current.right)
        if current.left:
            stack.append(current.left)


r = TreeNode(5)
t = r.left = TreeNode(2)
t.left = TreeNode(1)
t.right = TreeNode(3)
t = r.right = TreeNode(7)
t.left = TreeNode(6)
t.right = TreeNode(8)

preorder_stack(r)

输出:

5
2
1
3
7
6
8

2、中序遍历

算法:

  1. 如果栈顶存在左子节点,左子节点入栈
  2. 栈顶出栈,如果栈顶存在右子节点,右子节点入栈,如果不存在,继续出栈
  3. 重复第1步直到栈为空
from collections import deque


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

def inorder_stack(root: TreeNode):
    if not root:
        return
    stack = deque()
    stack.append(root)
    while stack:
        while stack[-1].left:
            stack.append(stack[-1].left)

        while stack:

            current = stack.pop()
            # 操作
            print(current.val)
            if current.right:
                stack.append(current.right)
                break


r = TreeNode(5)
t = r.left = TreeNode(2)
t.left = TreeNode(1)
t.right = TreeNode(3)
t = r.right = TreeNode(7)
t.left = TreeNode(6)
t.right = TreeNode(8)

inorder_stack(r)

输出:

1
2
3
5
6
7
8

3、后序遍历

后序遍历在中序的双层循环的基础上加入一个记录,专门记录上一次出栈的节点

算法:

  1. 如果栈顶存在左子节点,左子节点入栈
  2. 判断上一次出栈的节点是否是当前节点的右节点,如果是,当前节点出栈,如果不是并且存在右节点,将右节点入栈
  3. 重复第1步直到栈为空
from collections import deque

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

def postorder_stack(root: TreeNode):
    if not root:
        return
    stack = deque()
    last_node = None
    stack.append(root)
    while stack:
        while stack[-1].left:
            stack.append(stack[-1].left)

        while stack:
            if not stack[-1].right or last_node == stack[-1].right:
                current = stack.pop()
                # 操作
                last_node = current
                print(current.val)
            else:
                stack.append(stack[-1].right)
                break


r = TreeNode(5)
t = r.left = TreeNode(2)
t.left = TreeNode(1)
t.right = TreeNode(3)
t = r.right = TreeNode(7)
t.left = TreeNode(6)
t.right = TreeNode(8)

postorder_stack(r)

输出:

1
3
2
6
8
7
5