Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import sys
- from typing import List, Set
- class Node():
- def __init__(self, item):
- self.item = item
- self.parent = None
- self.left = None
- self.right = None
- self.color = 1
- class RedBlackTree():
- def __init__(self):
- self.TNULL = Node(0)
- self.TNULL.color = 0
- self.TNULL.left = None
- self.TNULL.right = None
- self.root = self.TNULL
- def pre_order_helper(self, node):
- if node != self.TNULL:
- sys.stdout.write(node.item + " ")
- self.pre_order_helper(node.left)
- self.pre_order_helper(node.right)
- def delete_fix(self, x):
- while x != self.root and x.color == 0:
- if x == x.parent.left:
- s = x.parent.right
- if s.color == 1:
- s.color = 0
- x.parent.color = 1
- self.left_rotate(x.parent)
- s = x.parent.right
- if s.left.color == 0 and s.right.color == 0:
- s.color = 1
- x = x.parent
- else:
- if s.right.color == 0:
- s.left.color = 0
- s.color = 1
- self.right_rotate(s)
- s = x.parent.right
- s.color = x.parent.color
- x.parent.color = 0
- s.right.color = 0
- self.left_rotate(x.parent)
- x = self.root
- else:
- s = x.parent.left
- if s.color == 1:
- s.color = 0
- x.parent.color = 1
- self.right_rotate(x.parent)
- s = x.parent.left
- if s.right.color == 0 and s.right.color == 0:
- s.color = 1
- x = x.parent
- else:
- if s.left.color == 0:
- s.right.color = 0
- s.color = 1
- self.left_rotate(s)
- s = x.parent.left
- s.color = x.parent.color
- x.parent.color = 0
- s.left.color = 0
- self.right_rotate(x.parent)
- x = self.root
- x.color = 0
- def __rb_transplant(self, u, v):
- if u.parent == None:
- self.root = v
- elif u == u.parent.left:
- u.parent.left = v
- else:
- u.parent.right = v
- v.parent = u.parent
- def delete_node_helper(self, node, key):
- z = self.TNULL
- while node != self.TNULL:
- if node.item == key:
- z = node
- if node.item <= key:
- node = node.right
- else:
- node = node.left
- if z == self.TNULL:
- return
- y = z
- y_original_color = y.color
- if z.left == self.TNULL:
- x = z.right
- self.__rb_transplant(z, z.right)
- elif (z.right == self.TNULL):
- x = z.left
- self.__rb_transplant(z, z.left)
- else:
- y = self.minimum(z.right)
- y_original_color = y.color
- x = y.right
- if y.parent == z:
- x.parent = y
- else:
- self.__rb_transplant(y, y.right)
- y.right = z.right
- y.right.parent = y
- self.__rb_transplant(z, y)
- y.left = z.left
- y.left.parent = y
- y.color = z.color
- if y_original_color == 0:
- self.delete_fix(x)
- def fix_insert(self, k):
- while k.parent.color == 1:
- if k.parent == k.parent.parent.right:
- u = k.parent.parent.left
- if u.color == 1:
- u.color = 0
- k.parent.color = 0
- k.parent.parent.color = 1
- k = k.parent.parent
- else:
- if k == k.parent.left:
- k = k.parent
- self.right_rotate(k)
- k.parent.color = 0
- k.parent.parent.color = 1
- self.left_rotate(k.parent.parent)
- else:
- u = k.parent.parent.right
- if u.color == 1:
- u.color = 0
- k.parent.color = 0
- k.parent.parent.color = 1
- k = k.parent.parent
- else:
- if k == k.parent.right:
- k = k.parent
- self.left_rotate(k)
- k.parent.color = 0
- k.parent.parent.color = 1
- self.right_rotate(k.parent.parent)
- if k == self.root:
- break
- self.root.color = 0
- def __print_helper(self, node, indent, last):
- if node != self.TNULL:
- sys.stdout.write(indent)
- if last:
- sys.stdout.write("R----")
- indent += " "
- else:
- sys.stdout.write("L----")
- indent += "| "
- s_color = "RED" if node.color == 1 else "BLACK"
- print(str(node.item) + "(" + s_color + ")")
- self.__print_helper(node.left, indent, False)
- self.__print_helper(node.right, indent, True)
- def preorder(self):
- self.pre_order_helper(self.root)
- def minimum(self, node):
- while node.left != self.TNULL:
- node = node.left
- return node
- def maximum(self, node):
- while node.right != self.TNULL:
- node = node.right
- return node
- def successor(self, x):
- if x.right != self.TNULL:
- return self.minimum(x.right)
- y = x.parent
- while y != self.TNULL and x == y.right:
- x = y
- y = y.parent
- return y
- def predecessor(self, x):
- if (x.left != self.TNULL):
- return self.maximum(x.left)
- y = x.parent
- while y != self.TNULL and x == y.left:
- x = y
- y = y.parent
- return y
- def left_rotate(self, x):
- y = x.right
- x.right = y.left
- if y.left != self.TNULL:
- y.left.parent = x
- y.parent = x.parent
- if x.parent == None:
- self.root = y
- elif x == x.parent.left:
- x.parent.left = y
- else:
- x.parent.right = y
- y.left = x
- x.parent = y
- def right_rotate(self, x):
- y = x.left
- x.left = y.right
- if y.right != self.TNULL:
- y.right.parent = x
- y.parent = x.parent
- if x.parent == None:
- self.root = y
- elif x == x.parent.right:
- x.parent.right = y
- else:
- x.parent.left = y
- y.right = x
- x.parent = y
- def insert(self, key):
- node = Node(key)
- node.parent = None
- node.item = key
- node.left = self.TNULL
- node.right = self.TNULL
- node.color = 1
- y = None
- x = self.root
- while x != self.TNULL:
- y = x
- if node.item < x.item:
- x = x.left
- else:
- x = x.right
- if y is not None and y.item.key == node.item.key:
- y.item.index.add(node.item.key)
- return
- node.parent = y
- if y == None:
- self.root = node
- elif node.item < y.item:
- y.left = node
- else:
- y.right = node
- if node.parent == None:
- node.color = 0
- return
- if node.parent.parent == None:
- return
- self.fix_insert(node)
- def get_root(self):
- return self.root
- def delete_node(self, item):
- self.delete_node_helper(self.root, item)
- def print_tree(self):
- self.__print_helper(self.root, "", True)
- class ServerStore:
- index: Set[int]
- key: int
- def __init__(self, key, ind):
- self.key, self.index = key, set()
- self.index.add(ind)
- def __lt__(self, other):
- return self.key < other.key
- def __gt__(self, other):
- return self.key > other.key
- def __le__(self, other):
- return self.key < other.key
- def __ge__(self, other):
- return self.key > other.key
- def __eq__(self, other):
- return len(self.index & other.index) > 0
- def __str__(self):
- return f"{self.index} {self.key}"
- n, m, q = map(int, input().split())
- bt = RedBlackTree()
- offs = [set() for _ in range(n)]
- off_count = [0] * n
- r = [0] * n
- for i in range(n):
- bt.insert(ServerStore(0, i))
- for _ in range(q):
- s = input()
- if s.startswith("DISABLE"):
- i, j = map(lambda x: int(x) - 1, s.split()[-2:])
- if j not in offs[i]:
- bt.delete_node_helper(bt.get_root(), ServerStore(r[i] * (m - off_count[i]), i))
- offs[i].add(j)
- off_count[i] += 1
- bt.insert(ServerStore(r[i] * (m - off_count[i]), i))
- elif s.startswith("RESET"):
- i = int(s.split()[-1]) - 1
- bt.delete_node_helper(bt.get_root(), ServerStore(r[i] * (m - off_count[i]), i))
- offs[i] = set()
- off_count[i] = 0
- r[i] += 1
- bt.insert(ServerStore(r[i] * (m - off_count[i]), i))
- elif s.startswith("GETMAX"):
- print(max(bt.maximum(bt.get_root()).item.index) + 1)
- else:
- print(min(bt.minimum(bt.get_root()).item.index) + 1)
Advertisement
Add Comment
Please, Sign In to add comment