searchtree.py

searchtree.py

# coding: utf-8
class Node:
	def __init__(self, data, parent):
		self.left = None
		self.data = data
		self.right = None
		self.parent = parent
		
	def __repr__(self):
		n = 'Node: '
		n += '{}'.format(self.data)
		
		if(self.left != None):
			n += ' left: {}'.format(self.left.data)
		if(self.right != None):
			n += ' right: {}'.format(self.right.data)
			
		return n
	
class Tree:
	def __init__(self):
		self.root = None
		
	def add(self, data):
		if(self.root == None):
			self.root = Node(data, None)
		else:	
			parent = self.find_parent(data)
			if(data <= parent.data):
				parent.left = Node(data, parent)
			else:
				parent.right = Node(data, parent)
				
	def __repr__(self):
		aslist = []
		if(self.root != None):
			aslist.append(self.root.data)
			self.walk_left(self.root.left, aslist)
			self.walk_right(self.root.right, aslist)
		
		return 'Tree: {}'.format(aslist)
		
	def walk_left(self, node, l):
		if(node != None):
			l.append(node.data)
			self.walk_left(node.left, l)
			self.walk_left(node.right, l)
		
	def walk_right(self, node, l):
		if(node != None):
			l.append(node.data)
			self.walk_right(node.left, l)
			self.walk_right(node.right, l)
		
		
	def find_parent(self, data):
		
		current = self.root
		
		while(current != None):
			parent = current
			if(data <= current.data):
				current = current.left
			else:
				current = current.right
				
		return parent
		
	def find(self, data):
	
		current = self.root
		
		while(current != None):
			
			if(data==current.data):
				return current
			if(data<current.data):
				current = current.left
			else:
				current = current.right
				
		return current
				
				
t = Tree()		
t.add(5)
t.add(2)
t.add(3)
t.add(1)
t.add(8)
t.add(6)
		
print(t)

n = t.find(2)

print("find node for 2 and print left and right")
print(n)