Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- class AVLTree:
- def insert(self, root, key, value):
- #Normal BST Insertion
- if root is None:
- return AVLTreeNode(key, value)
- elif key < root.key:
- root.left = self.insert(root.left, key, value)
- elif key > root.key:
- root.right = self.insert(root.right, key, value)
- else:
- root.values.insert(0, value)
- #Update Height
- root.height = 1+max(self.height(root.left), self.height(root.right))
- #Check balance
- balance = self.balance(root)
- #Rotate if needed
- #Left-Left
- if balance > 1 and key < root.left.key:
- return self.rRotate(root)
- #Right Right
- if balance < -1 and key > root.right.key:
- return self.lRotate(root)
- #Left Right
- if balance > 1 and key > root.left.key:
- root.left = self.lRotate(root.left)
- return self.rRotate(root)
- #Right Left
- if balance < -1 and key < root.right.key:
- root.right = self.rRotate(root.right)
- return self.lRotate(root)
- return root
- def inside(self, root, key, value):
- if root is None:
- return False
- elif key == root.key:
- return value in root.values
- elif key < root.key:
- return self.search(root.left, key, value)
- else:
- return self.search(root.right, key, value)
- def search(self, root, key):
- if root is None:
- return None
- elif key == root.key:
- return root
- elif key < root.key:
- return self.search(root.left, key)
- else:
- return self.search(root.right, key)
- def lessThan(self, root, high, equal):
- if root is None:
- return []
- r = []
- if root.left is not None:
- r += self.lessThan(root.left, high, equal)
- if root.key < high or (equal and root.key == high):
- r.append(root)
- if root.right is not None:
- r += self.lessThan(root.right, high, equal)
- return r
- def greaterThan(self, root, low, equal):
- if root is None:
- return []
- r = []
- if root.left is not None:
- r += self.lessThan(root.left, high, equal)
- if root.key > low or (equal and root.key == low):
- r.append(root)
- if root.right is not None:
- r += self.lessThan(root.right, high, equal)
- return r
- def inRange(self, root, low, high, equal1, equal2):
- if root is None:
- return []
- r = []
- if root.left is not None:
- r += self.inRange(root.left, low, high, equal1, equal2)
- if root.key < high and root.key > low or (equal2 and root.key == high) or (equal1 and root.key == low):
- r.append(root)
- if root.right is not None:
- r += self.inRange(root.right, low, high, equal1, equal2)
- return r
- def deleteLessThan(self, root, high, equal):
- if root is None:
- return
- if root.left is not None:
- root.left = self.deleteLessThan(root.left, high, equal)
- if root.right is not None:
- root.right = self.deleteLessThan(root.right, high, equal)
- if root.key < high or (equal and root.key == high):
- root = self.delete(root, root.key, None)
- return root
- def deleteGreaterThan(self, root, low, equal):
- if root is None:
- return
- if root.left is not None:
- root.left = self.deleteGreaterThan(root.left, low, equal)
- if root.right is not None:
- root.right = self.deleteGreaterThan(root.right, low, equal)
- if root.key > low or (equal and root.key == low):
- root = self.delete(root, root.key, None)
- return root
- def deleteInRange(self, root, low, high, equal1, equal2):
- return root
- def delete(self, root, key, value):
- #Standard BST Delete
- if root is None:
- return root
- elif key < root.key:
- root.left = self.delete(root.left, key, value)
- elif key > root.key:
- root.right = self.delete(root.right, key, value)
- else:
- if value != None and len(root.values) > 1:
- root.values.remove(value)
- return root
- if root.left is None:
- temp = root.right
- root = None
- return temp
- elif root.right is None:
- temp = root.left
- root = None
- return temp
- temp = self.minimum(root.right)
- root.key = temp.key
- root.values = temp.values
- root.right = self.delete(root.right, temp.key,value)
- #If it is gone
- if root is None:
- return root
- #Update Height
- root.height = 1+max(self.height(root.left), self.height(root.right))
- #Check balance
- balance = self.balance(root)
- #Rotate if needed
- #Left-Left
- if balance > 1 and self.balance(root.left) >= 0:
- return self.rRotate(root)
- #Right Right
- if balance < -1 and self.balance(root.right) <= 0:
- return self.lRotate(root)
- #Left Right
- if balance > 1 and self.balance(root.left) < 0:
- root.left = self.lRotate(root.left)
- return self.rRotate(root)
- #Right Left
- if balance < -1 and self.balance(root.right) > 0:
- root.right = self.rRotate(root.right)
- return self.lRotate(root)
- return root
- def height(self, root):
- #If it's not a tree it doesn't have a height
- #Else just get the height
- if root is None: return 0
- return root.height
- def balance(self, root):
- #If it's not a tree it doesn't have a balance
- #Else get the difference between height of subtrees
- if root is None: return 0;
- return self.height(root.left)-self.height(root.right)
- def lRotate(self, root):
- #Temporary Storage
- r = root.right
- rl = r.left
- #Swap Around
- r.left = root
- root.right = rl
- #Update Heights
- root.height = 1+max(self.height(root.left), self.height(root.right))
- r.height = 1+max(self.height(r.left), self.height(r.right))
- return r
- def rRotate(self, root):
- #Temporary Storage
- l = root.left
- lr = l.right
- #Swap Around
- l.right = root
- root.left = lr
- #Update Heights
- root.height = 1+max(self.height(root.left), self.height(root.right))
- l.height = 1+max(self.height(l.left), self.height(l.right))
- return l
- def minimum(self, root):
- #If we reach the end just return
- #Else continue going left
- if root is None or root.left is None: return root
- return self.minimum(root.left)
- def string(self, root):
- #If it isn't a node represent it as a blank
- #Else return the left child, itself, and the right child
- if root is None: return "(-)"
- return "("+self.string(root.left)+str(root.key)+self.string(root.right)+")"
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement