Binary search tree traversal and depth first search in Python

Question

I am learning about the tree traversal, and I would like to ask for feedback.

Using recursion, write a function that takes in a binary search tree and number of outputs a list of values ordered by depth first traversal, pre-order, in-order, post-order.

#    PRE-ORDER
#
#      Input: node {TreeNode}
#      Output: {List}
#
#      NOTE: Confirm with your answer from problem 2b.
#


def dfs_pre(node):
    result = []
    def traverse(node):
      nonlocal result 
      if node is None: return 
      result.append(node.value)
      traverse(node.left)
      traverse(node.right)

    traverse(node)
    return result

#
#  Using recursion, write a function that takes in a binary search tree and
#      outputs a list of values ordered by IN-ORDER DEPTH FIRST traversal.
#
#      Input: node {TreeNode}
#      Output: {List}
#
#      NOTE: Confirm with your answer from problem 2c.
#


def dfs_in(node):
    # YOUR WORK HERE
    result = []
    def traverse(node):
      nonlocal result 
      if node is None: return 
      traverse(node.left)
      result.append(node.value)
      traverse(node.right)

    traverse(node)
    return result

#
#   Using recursion, write a function that takes in a binary search tree and
#      outputs a list of values ordered by POST-ORDER DEPTH FIRST traversal.
#
#      Input: Binary Search Tree
#      Output: List
#
#      NOTE: Confirm with your answer from problem 2d.
#


def dfs_post(node):
    # YOUR WORK HERE
    result = []
    def traverse(node):
      nonlocal result 
      if node is None: return 
      traverse(node.left)
      traverse(node.right)
      result.append(node.value)

    traverse(node)
    return result

Here's the basic testing:

from io import StringIO
import sys


class Capturing(list):
    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        return self

    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        sys.stdout = self._stdout

# code for checking if lists are equal
def lists_equal(lst1, lst2):
    if len(lst1) != len(lst2):
        return False
    for i in range(0, len(lst1)):
        if lst1[i] != lst2[i]:
            return False
    return True


# generate test tree for the rest of the tests
test_tree = TreeNode(4)
test_tree.left = TreeNode(2)
test_tree.left.left = TreeNode(1)
test_tree.left.right = TreeNode(3)
test_tree.right = TreeNode(5)
test_tree.right.right = TreeNode(7)
test_tree.right.right.left = TreeNode(6)
test_tree.right.right.right = TreeNode(8)



print('pre-order depth first search tests')
test_count = [0, 0]


def test():
    results = dfs_pre(test_tree)
    return lists_equal(results, [4, 2, 1, 3, 5, 7, 6, 8])


expect(test_count, 'able to return values of BST in pre-order depth first manner - [4,2,1,3,5,7,6,8]', test)


def test():
    results = dfs_pre(deserialize([]))
    return lists_equal(results, [])


expect(test_count, 'returns an empty erray for an empty BST', test)

print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + '\n\n')


print('in-order depth first search tests')
test_count = [0, 0]


def test():
    results = dfs_in(test_tree)
    return lists_equal(results, [1, 2, 3, 4, 5, 6, 7, 8])


expect(test_count, 'able to return values of BST in an in-order depth first manner - [1,2,3,4,5,6,7,8]', test)


def test():
    results = dfs_in(deserialize([]))
    return lists_equal(results, [])


expect(test_count, 'returns an empty erray for an empty BST', test)

print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + '\n\n')


print('post-order depth first search tests')
test_count = [0, 0]


def test():
    results = dfs_post(test_tree)
    return lists_equal(results, [1, 3, 2, 6, 8, 7, 5, 4])


expect(test_count, 'able to return values of BST in post-order depth first manner - [1,3,2,6,8,7,5,4]', test)


def test():
    results = dfs_post(deserialize([]))
    return lists_equal(results, [])


expect(test_count, 'returns an empty erray for an empty BST', test)

print('PASSED: ' + str(test_count[0]) + ' / ' + str(test_count[1]) + '\n\n')

It passed all of the tests.

pre-order depth first search tests
  1)    true : able to return values of BST in pre-order depth first manner - [4
,2,1,3,5,7,6,8]
  2)    true : returns an empty erray for an empty BST
PASSED: 2 / 2


in-order depth first search tests
  1)    true : able to return values of BST in an in-order depth first manner -
[1,2,3,4,5,6,7,8]
  2)    true : returns an empty erray for an empty BST
PASSED: 2 / 2


post-order depth first search tests
  1)    true : able to return values of BST in post-order depth first manner - [
1,3,2,6,8,7,5,4]
  2)    true : returns an empty erray for an empty BST
PASSED: 2 / 2

Answer 1

You should probably replace your # YOUR WORK HERE comments with """docstrings""" describing what the functions do.

---

Your lists_equal() function could be made shorter, and perhaps a little clearer, by using all(), zip() and list comprehension:

def lists_equal(lst1, lst2):
    if len(lst1) != len(lst2):
        return False
    return all(item1 == item2 for item1, item2 in zip(lst1, lst2))