Cracking the coding interview in python – 4.4 a tree traversal

It is a rainy day. Here in Nice, as everywhere, people use their cars as umbrellas so the traffic jams are bigger that usual. This gives me the time to solve another exercise from the Cracking the Coding interview.

Here it is:

4.4 Given a binary search tree, design an algorithm which creates a linked list of all the
nodes at each depth (eg, if you have a tree with depth D, you’ll have D linked lists).

Well, nothing special about this. Just a tree traversal and some lists creation. Let’s try to do something nice. I will first write my tree class and implement a visitor pattern. The visitor will be a function. Then, I will write a small function to create and fill the linked lists.

First, lets write the TreeNode:


class TreeNode:
	
		
	def __init__(self, data, left_node=None, right_node=None):
		self.data=data
		self.left=left_node
		self.right=right_node
			
	#in-order: left, current, right
	def traverse(self,func, level=0):
		if self.left:
			self.left.traverse(func, level+1)
		func(self,level)	
		if self.right:
			self.right.traverse(func, level+1)

Python allows me to pass a function as argument, this makes the visitor pattern look nice and simple.
Now, if I want to print the tree, I would write a print function:

def printNode(node, level):
	print level*'   '+`node.data`

I could use the function from my previous post to create a tree. Then I’ll print it:

def printNode(node, level):
	print level*'   '+`node.data`

def test():	
	arr=[1,2,3,4,5,7,9,11,66]
	t=create_tree(arr)
	t.traverse(printNode)

And now, lets solve today’s exercise and create those lists. I’ll use a dictionnary as support for the lists. The key is the level, the value is the link list.

	d = dict()
	def fillDict(node, level):
		if d.has_key(level):
			l = d[level]
			l.append(node)
		else:
			l = [node]
			d[level]=l
	treeNode.traverse(fillDict)
	return  d

So what do I do ? I create a dictionary d.
Then, I write my function (fillDict) which takes in argument a node and a level and puts it at the right position in the dictionary.
Then I pass the function as argument to the traverse function.
What is noticeable here is that the fillDict function, executed in the scope of traverse function, still have accesses to its own enclosing scope, where d is located.

Update 1 – let’s make it more pythonic

mebubo hates my code. He wants to do it like this:

def createListsFromTree2(treeNode):
	d = {}
	def fillDict(node, level):
		d.setdefault(level, []).append(node)
	treeNode.traverse(fillDict)
	return  d

Here we don’t need all the if that I do this else I do that. We say that the defaut value for d[level] is [], to be used for first initialization.

Update 2 – let’s make it (even) more pythonic

mebubo wants to do even better:

from collections import defaultdict
def createListsFromTree3(treeNode):
	d = defaultdict(list)
	def fillDict(node, level):
		d[level].append(node)
	treeNode.traverse(fillDict)
	return  d

That’s all folks.
The rained stopped.

Advertisements
This entry was posted in Cracking the coding interview - a python experience and tagged , , , , . Bookmark the permalink.

One Response to Cracking the coding interview in python – 4.4 a tree traversal

  1. mebubo says:

    We could also use lambda, but I’m not sure if this version is more readable or not:

    from collections import defaultdict
    def createListsFromTree3(treeNode):
    	d = defaultdict(list)
    	treeNode.traverse(lambda node, level: d[level].append(node))
    	return  d
    

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s