diff --git a/contents/tree_traversal/code/python/tree_traversal.py b/contents/tree_traversal/code/python/tree_traversal.py
index 735837bdd..1c15eae23 100644
--- a/contents/tree_traversal/code/python/tree_traversal.py
+++ b/contents/tree_traversal/code/python/tree_traversal.py
@@ -46,31 +46,18 @@ def dfs_recursive_inorder_btree(node):
def dfs_stack(node):
- stack = []
- stack.append(node)
-
- temp = None
-
- while len(stack) > 0:
- print(stack[-1].data, end=' ')
- temp = stack.pop()
-
- for child in temp.children:
- stack.append(child)
-
+ stack = [node]
+ while stack:
+ node = stack.pop()
+ stack.extend(node.children)
+ print(node.data, end=' ')
def bfs_queue(node):
- queue = []
- queue.append(node)
-
- temp = None
-
- while len(queue) > 0:
- print(queue[0].data, end=' ')
- temp = queue.pop(0)
-
- for child in temp.children:
- queue.append(child)
+ queue = [node]
+ while queue:
+ node = queue.pop(0)
+ queue.extend(node.children)
+ print(node.data)
def main():
diff --git a/contents/tree_traversal/tree_traversal.md b/contents/tree_traversal/tree_traversal.md
index 566aca469..2defe9da8 100644
--- a/contents/tree_traversal/tree_traversal.md
+++ b/contents/tree_traversal/tree_traversal.md
@@ -231,7 +231,7 @@ In code, it looks like this:
{% sample lang="js" %}
[import:53-60, lang:"javascript"](code/javascript/tree.js)
{% sample lang="py" %}
-[import:48-59, lang:"python"](code/python/tree_traversal.py)
+[import:48-53, lang:"python"](code/python/tree_traversal.py)
{% sample lang="scratch" %}
@@ -284,7 +284,7 @@ And this is exactly what Breadth-First Search (BFS) does! On top of that, it can
{% sample lang="js" %}
[import:62-69, lang:"javascript"](code/javascript/tree.js)
{% sample lang="py" %}
-[import:62-72, lang:"python"](code/python/tree_traversal.py)
+[import:55-60, lang:"python"](code/python/tree_traversal.py)
{% sample lang="scratch" %}
