优化下面代码class TreeNode: def __init__(self, val): self.val = val self.left = None self.right = None root = TreeNode('a') root.left = TreeNode('b') root.right = TreeNode('c') root.left.left = TreeNode('d') root.left.right = TreeNode('e') root.right.left = TreeNode('f') root.right.right = TreeNode('g') root.left.left.left = TreeNode('h') root.left.left.right = TreeNode('i') def preorder_traversal(root): if not root: return print(root.val, end=' ') preorder_traversal(root.left) preorder_traversal(root.right) def inorder_traversal(root): if not root: return inorder_traversal(root.left) print(root.val, end=' ') inorder_traversal(root.right) def postorder_traversal(root): if not root: return postorder_traversal(root.left) postorder_traversal(root.right) print(root.val, end=' ') from collections import deque def level_order_traversal(root): if not root: return queue = deque() queue.append(root) while queue: node = queue.popleft() print(node.val, end=' ') if node.left: queue.append(node.left) if node.right: queue.append(node.right) def get_height(root): if not root: return 0 left_height = get_height(root.left) right_height = get_height(root.right) return max(left_height, right_height) + 1 def get_node_count(root): if not root: return 0 left_node_count = get_node_count(root.left) right_node_count = get_node_count(root.right) return left_node_count + right_node_count + 1 print("先序遍历:") preorder_traversal(root) print("中序遍历:") inorder_traversal(root) print("后序遍历:") postorder_traversal(root) print("层次遍历:") level_order_traversal(root) print("该二叉树的高度为:") get_height(root) print("该二叉树的节点个数为 ") get_node_count(root)
时间: 2024-02-25 11:59:54 浏览: 57
这段代码实现了二叉树的四种遍历方式和获取二叉树的高度和节点个数。代码本身没有什么问题,但是可以考虑进行一些优化。
1. 可以在遍历时使用yield关键字,将遍历结果逐个返回,而不是一次性打印出来。这样可以节省内存空间,避免一次性将所有结果存储在内存中。
2. 可以使用迭代方式代替递归方式实现遍历,这样可以避免递归深度过大导致的栈溢出问题。
3. 可以将获取二叉树高度和节点个数的函数合并成一个函数,减少重复代码。
下面是优化后的代码示例:
```python
class TreeNode:
def __init__(self, val):
self.val = val
self.left = None
self.right = None
def preorder_traversal(root):
if not root:
return
yield root.val
yield from preorder_traversal(root.left)
yield from preorder_traversal(root.right)
def inorder_traversal(root):
if not root:
return
yield from inorder_traversal(root.left)
yield root.val
yield from inorder_traversal(root.right)
def postorder_traversal(root):
if not root:
return
yield from postorder_traversal(root.left)
yield from postorder_traversal(root.right)
yield root.val
def level_order_traversal(root):
if not root:
return
queue = [root]
while queue:
node = queue.pop(0)
yield node.val
if node.left:
queue.append(node.left)
if node.right:
queue.append(node.right)
def get_tree_info(root):
if not root:
return 0, 0
left_height, left_node_count = get_tree_info(root.left)
right_height, right_node_count = get_tree_info(root.right)
return max(left_height, right_height) + 1, left_node_count + right_node_count + 1
root = TreeNode('a')
root.left = TreeNode('b')
root.right = TreeNode('c')
root.left.left = TreeNode('d')
root.left.right = TreeNode('e')
root.right.left = TreeNode('f')
root.right.right = TreeNode('g')
root.left.left.left = TreeNode('h')
root.left.left.right = TreeNode('i')
print("先序遍历:", list(preorder_traversal(root)))
print("中序遍历:", list(inorder_traversal(root)))
print("后序遍历:", list(postorder_traversal(root)))
print("层次遍历:", list(level_order_traversal(root)))
height, node_count = get_tree_info(root)
print("该二叉树的高度为:", height)
print("该二叉树的节点个数为:", node_count)
```