Updated tree traversal chapter (algorithm-archivists#979)

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():

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" %}
 <p>
   <img  class="center" src="code/scratch/dfs-stack.svg" style="width:70%" />
@@ -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" %}
 <p>
   <img  class="center" src="code/scratch/bfs.svg" style="width:70%" />