From 50436468d33f5d74b56820c120fad5098c6e26df Mon Sep 17 00:00:00 2001 From: marina <138340846+bt3gl-cryptographer@users.noreply.github.com> Date: Thu, 3 Aug 2023 15:37:11 -0700 Subject: [PATCH] add avl tree class --- trees/bst_avl.py | 155 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 trees/bst_avl.py diff --git a/trees/bst_avl.py b/trees/bst_avl.py new file mode 100644 index 0000000..134587f --- /dev/null +++ b/trees/bst_avl.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# author: bt3gl + + +class Node(): + def __init__(self, val): + self.val = val + self.left = None + self.right = None + self.height = 1 + + +class AVLTree(): + + def __init__(self): + self.root = None + self.size = 0 + + def height(self, node): + if node: + return node.height + return 0 + + def set_height(self, node): + if node is None: + return 0 + return 1 + max(self.height(node.left), self.height(node.right)) + + def right_rotate(self, node): + new_root = node.left + node.left = node.left.right + new_root.right = node + node.height = self.set_height(node) + new_root.height = self.set_height(new_root) + return new_root + + def left_rotate(self, node): + new_root = node.right + node.right = node.right.left + new_root.left = node + node.height = self.set_height(node) + new_root.height = self.set_height(new_root) + return new_root + + def insert(self, node, val): + if node == self.root: + self.size += 1 + if node is None: + return Node(val) + if node.val < val: + node.right = self.insert(node.right, val) + else: + node.left = self.insert(node.left, val) + balance = self.height(node.left) - self.height(node.right) + if balance > 1: + if self.height(node.left.left) > self.height(node.left.right): + node = self.right_rotate(node) + else: + node.left = self.left_rotate(node.left) + node = self.right_rotate(node) + elif balance < -1: + if self.height(node.right.right) > self.height(node.right.left): + node = self.left_rotate(node) + else: + node.right = self.right_rotate(node.right) + node = self.left_rotate(node) + else: + node.height = self.set_height(node) + return node + + def get_min_val(self, node): + if node is None or node.left is None: + return node + return self.get_min_val(node.left) + + def remove(self, node, val): + if node is None: + return node + if node.val < val: + node.right = self.remove(node.right, val) + elif node.val > val: + node.left = self.remove(node.left, val) + else: + if node.left is None: + return node.right + elif node.right is None: + return node.left + else: + right_min_val_node = self.get_min_val(node.right) + node.val = right_min_val_node.val + node.right = self.remove(node.right, right_min_val_node.val) + + node.height = self.set_height(node) + balance = self.height(node.left) - self.height(node.right) + if balance > 1: + if self.height(node.left.left) > self.height(node.left.right): + node = self.right_rotate(node) + else: + node.left = self.left_rotate(node.left) + node = self.right_rotate(node) + elif balance < -1: + if self.height(node.right.right) > self.height(node.right.left): + node = self.left_rotate(node) + else: + node.right = self.right_rotate(node.right) + node = self.left_rotate(node) + else: + node.height = self.set_height(node) + return node + + def predecessor(self, node, val): + if node is None: + return node + if node.val == val: + return val + elif node.val > val: + return self.predecessor(node.left, val) + else: + right_res = self.predecessor(node.right, val) + return right_res if right_res else node.val + + def successor(self, node, val): + if node is None: + return node + if node.val == val: + return val + elif node.val < val: + return self.successor(node.right, val) + else: + left_res = self.successor(node.left, val) + return left_res if left_res else node.val + + +def contains_duplicate_near(self, nums, k, t): + + avltree = AVLTree() + root = avltree.root + + + for i, num in enumerate(nums): + predecessor = avltree.predecessor(root, num) + if predecessor is not None and abs(predecessor - num) <= t: + return True + successor = avltree.successor(root, num) + if successor is not None and abs(successor - num) <= t: + return True + + root = avltree.insert(root, num) + + if avltree.size > k: + root = avltree.remove(root, nums[i-k]) + + return False +