InOrderIterator now returns the node_data instead of the subtree. Fixes

#3.
skytreader · Mar 27, 2015 · 643c414 · 643c414
1 parent 8c89836
commit 643c414
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/data_structures/binary_tree.py b/data_structures/binary_tree.py
@@ -127,7 +127,7 @@ def __next__(self):
         # been visited.
 
         self.visited.append(next_node)
-        return next_node
+        return next_node.node_data
 
 class NaiveBinaryTree(BinaryTree):
     """

diff --git a/data_structures/tests/binary_tree_test.py b/data_structures/tests/binary_tree_test.py
@@ -75,8 +75,8 @@ def test_dfs(self):
         dfs = InorderIterator(self.test_root)
         iterator_order = ""
 
-        for node in dfs:
-            iterator_order += node.node_data
+        for data in dfs:
+            iterator_order += data
 
         self.assertEqual(self.dfs_order, iterator_order)
 
@@ -93,8 +93,8 @@ def test_dfs(self):
         dfs = InorderIterator(self.boole_tree)
         iterator_order = []
 
-        for node in dfs:
-            iterator_order.append(node.node_data)
+        for data in dfs:
+            iterator_order.append(data)
 
         self.assertEqual([True, False, True, True, False], iterator_order)
 
@@ -111,8 +111,8 @@ def test_dfs(self):
         dfs = InorderIterator(self.boole_tree)
         iterator_order = []
 
-        for node in dfs:
-            iterator_order.append(node.node_data)
+        for data in dfs:
+            iterator_order.append(data)
 
         self.assertEqual(["D", "B", "E", "A", "C"], iterator_order)