Breadth-First Search

bfs.py

"""
Breadth-First Search for binary tree
"""

import unittest

class Node(object):
    def __init__(self, data = None, left = None, right = None):
        self.data = data
        self.left = left
        self.right = right

    def setLeft(self, node = None):
        self.left = node

    def setRight(self, node = None):
        self.right = node


class Queue:
    def __init__(self):
        self.items = []

    def isEmpty(self):
        return self.items == []

    def enqueue(self, item):
        self.items.insert(0, item)

    def dequeue(self):
        return self.items.pop()



def BFS(root):
    q = Queue()
    q.enqueue(root)
    visited = []
    res = [] # list with nodes data

    if root == None:
        return res

    while not q.isEmpty():
        node = q.dequeue()
        visited.append(node)

        if node.left:
            q.enqueue(node.left)
        if node.right:
            q.enqueue(node.right)
    
    # return visited
    for i in visited:
        res.append(i.data)
    return res


class TestMethods(unittest.TestCase):

    def test_bfs(self):
        self.assertEqual(BFS(root), [6, 4, 8, 3, 5, 7, 9, 2])

    def test_empty_bfs(self):
        self.assertEqual(BFS(None), [])

    def test_single_item_bfs(self):
        self.assertEqual(BFS(root1), [500])    



n2 = Node(2)
n3 = Node(3, n2)
n5 = Node(5)
n4 = Node(4, n3, n5)
n7 = Node(7)
n9 = Node(9)
n8 = Node(8, n7, n9)
root = Node(6, n4, n8)

# another tree
root1 = Node(500)

print BFS(root)


unittest.main()