-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathtraversal_pre_order.py
138 lines (122 loc) · 4.98 KB
/
traversal_pre_order.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from .binary_tree import TreeNode
import unittest
import collections
from typing import List
# 迭代版的前序和中序遍历用的不是一套模板
def in_order_iterative(root: TreeNode) -> List[int]:
result = []
if root is None:
return result
# stack表示从根节点到当前节点的路径上的所有节点
stack = collections.deque()
curr_node = root
while curr_node or stack:
while curr_node:
stack.append(curr_node)
curr_node = curr_node.left
curr_node = stack.pop()
result.append(curr_node.val)
curr_node = curr_node.right
return result
# 230. Kth Smallest Element in a BST
def bst_kth_smallest(root: TreeNode, k: int) -> int:
"""
TODO 优化方法
用一个HashMap<TreeNode, Integer>记录每个节点的子树的节点个数
利用快速选择的算法,确定第k小的数在哪一个子树中,时间复杂度为O(h)
例如求k=20小的元素,root.left有15个,root算一个,那么等同于求root.right中第4小的元素
"""
stack = collections.deque()
curr_node = root
while curr_node or stack:
while curr_node:
stack.append(curr_node)
curr_node = curr_node.left
curr_node = stack.pop()
k -= 1
if k == 0:
return curr_node.val
curr_node = curr_node.right
# 正统的后续遍历思维的迭代法,基于中序遍历版本稍作修改
def post_order_iterative_2(root: TreeNode) -> List[int]:
result = []
if root is None:
return result
stack = collections.deque()
curr_node = root
while curr_node or stack:
while curr_node:
stack.append(curr_node)
# 判断当前节点的左子树是否存在,若存在则持续往深处遍历左子树,若不存在就转向右子树
if curr_node.left:
curr_node = curr_node.left
else:
curr_node = curr_node.right
# 走到这里时只可能是最底层且未记录的节点
curr_node = stack.pop()
result.append(curr_node.val)
# 若栈不为空且当前节点是栈顶元素的左节点,说明已记录栈顶元素左节点的值
if stack and stack[-1].left is curr_node:
# 则转为遍历栈顶节点的右子树
curr_node = stack[-1].right
else:
# 栈顶元素没有左子树和右子树,则curr_node设为None让栈顶元素出栈
curr_node = None
return result
def post_order_iterative_3(root: TreeNode) -> List[int]:
"""
1、如果根节点非空,将根节点加入到栈中。
2、如果栈不空,取栈顶元素(暂时不弹出),
a.如果(左子树已访问过或者左子树为空),且(右子树已访问过或右子树为空),则弹出栈顶节点,将其值加入数组,
b.如果左子树不为空,且未访问过,则将左子节点加入栈中,并标左子树已访问过。
c.如果右子树不为空,且未访问过,则将右子节点加入栈中,并标右子树已访问过。
3、重复第二步,直到栈空。
"""
result = []
stack = []
prev, curr = None, root
if not root:
return result
# 通过维护curr前继节点的后序遍历
stack.append(root)
while len(stack) > 0:
curr = stack[-1]
if not prev or prev.left == curr or prev.right == curr:
# traverse down the tree
if curr.left:
stack.append(curr.left)
elif curr.right:
stack.append(curr.right)
elif curr.left == prev:
# traverse up the tree from the left
if curr.right:
stack.append(curr.right)
else:
# traverse up the tree from the right
result.append(curr.val)
stack.pop()
prev = curr
return result
class Testing(unittest.TestCase):
TEST_CASES = [
{'binary_tree': "1()(2(3))", 'pre_order': [1, 2, 3], 'in_order': [1, 3, 2], 'post_order': [3, 2, 1]},
{'binary_tree': "3(9)(20(15)(7))",
'pre_order': [3, 9, 20, 15, 7],
'in_order': [9, 3, 15, 20, 7],
'post_order': [9, 15, 7, 20, 3]},
]
def test_pre_order_iterative(self):
for data in self.TEST_CASES:
root = TreeNode.from_str(data['binary_tree'])
pre_order = pre_order_iterative(root)
self.assertEqual(data['pre_order'], pre_order)
def test_in_order_iterative(self):
for data in self.TEST_CASES:
root = TreeNode.from_str(data['binary_tree'])
pre_order = in_order_iterative(root)
self.assertEqual(data['in_order'], pre_order)
def test_post_order_iterative(self):
for data in self.TEST_CASES:
root = TreeNode.from_str(data['binary_tree'])
self.assertEqual(data['post_order'], post_order_iterative_1(root))
self.assertEqual(data['post_order'], post_order_iterative_2(root))