from binarytree import * class BinarySearchTree(BinaryTree): def __init__(self, key, value): BinaryTree.__init__(self, value) self.key = key # Puts the key-value pair into the tree. # If the key is already in the tree, then changes its value. # Returns 0 if key was already in tree, 1 if it was not. # That is, the return value is the net change in the number of nodes. def put(self, key, value): if key < self.key: if self.leftChild == None: self.leftChild = BinarySearchTree(key, value) return 1 else: return self.leftChild.put(key, value) elif key > self.key: if self.rightChild == None: self.rightChild = BinarySearchTree(key, value) return 1 else: return self.rightChild.put(key, value) else: self.setRootValue(value) return 0 # Returns the value associated to the given key, or None. def get(self, key): if key < self.key: if self.leftChild == None: return None else: return self.leftChild.get(key) elif key > self.key: if self.rightChild == None: return None else: return self.rightChild.get(key) else: # The key must equal the key stored here. return self.value # Deletes the least key from the tree. # Returns a pair [root, node] where root is the new root of the tree # and node is the tree node containing the deleted key. # If one deletes the last key in the tree, then root == None. def deleteLeastKey(self): if self.leftChild == None: # This node is the least key node. return [self.rightChild, self] else: # The least key node is somewhere in the left subtree. [root, node] = self.leftChild.deleteLeastKey() self.leftChild = root return [self, node] # Deletes the key from the tree. # Returns a pair [root, node] where root is the new root of the tree # and node is the tree node containing the deleted key. # If the key is not in the tree, then root == tree and node == None. # If one deletes the last key in the tree, then root == None. def delete(self, key): if key < self.key: if self.leftChild == None: # Key not found. return [self, None] else: # Key may be in left subtree. [root, node] = self.leftChild.delete(key) self.leftChild = root return [self, node] elif key > self.key: if self.rightChild == None: # Key not found. return [self, None] else: # Key may be in right subtree. [root, node] = self.rightChild.delete(key) self.rightChild = root return [self, node] else: # The key must equal the key stored here. if self.leftChild == None: if self.rightChild == None: # There are no children. return [None, self] else: # There is a right child but no left child. return [self.rightChild, self] else: if self.rightChild == None: # There is a left child but no right child. return [self.leftChild, self] else: # There are two children; this is the difficult case. [root, node] = self.rightChild.deleteLeastKey() node.setLeftChild(self.leftChild) node.setRightChild(root) return [node, self] def __str__(self): # This returns a tree-like string. # return '''[(''' + str(self.key) + ''', ''' + str(self.value) + '''), ''' + str(self.leftChild) + ''', ''' + str(self.rightChild) + ''']''' # This other code returns a dictionary-like string. string = '''{''' + str(self.key) + ''':''' + str(self.value) if self.leftChild != None: string = string + ''', ''' + str(self.leftChild) if self.rightChild != None: # Concatenate string for right subtree, stripping off braces. string = string + ''', ''' + str(self.rightChild) string = string + '''}''' return string # If the user ran this file explicitly, then run this demo code. if __name__ == "__main__": tree = BinarySearchTree(13, 'dog') print str(tree) tree.put(17, 'jane') print str(tree) tree.put(23, 'juan') print str(tree) tree.put(14, 'john') print str(tree) tree.put(20, 'ian') print str(tree) tree.put(28, 'joan') print str(tree) tree.put(19, 'jan') print str(tree) tree.delete(23) print str(tree) print str(tree.height())