Leamich

контест шбр а

Apr 11th, 2023
668
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 9.71 KB | None | 0 0
  1. import sys
  2. from typing import List, Set
  3.  
  4.  
  5. class Node():
  6.     def __init__(self, item):
  7.         self.item = item
  8.         self.parent = None
  9.         self.left = None
  10.         self.right = None
  11.         self.color = 1
  12.  
  13.  
  14. class RedBlackTree():
  15.     def __init__(self):
  16.         self.TNULL = Node(0)
  17.         self.TNULL.color = 0
  18.         self.TNULL.left = None
  19.         self.TNULL.right = None
  20.         self.root = self.TNULL
  21.  
  22.     def pre_order_helper(self, node):
  23.         if node != self.TNULL:
  24.             sys.stdout.write(node.item + " ")
  25.             self.pre_order_helper(node.left)
  26.             self.pre_order_helper(node.right)
  27.  
  28.     def delete_fix(self, x):
  29.         while x != self.root and x.color == 0:
  30.             if x == x.parent.left:
  31.                 s = x.parent.right
  32.                 if s.color == 1:
  33.                     s.color = 0
  34.                     x.parent.color = 1
  35.                     self.left_rotate(x.parent)
  36.                     s = x.parent.right
  37.  
  38.                 if s.left.color == 0 and s.right.color == 0:
  39.                     s.color = 1
  40.                     x = x.parent
  41.                 else:
  42.                     if s.right.color == 0:
  43.                         s.left.color = 0
  44.                         s.color = 1
  45.                         self.right_rotate(s)
  46.                         s = x.parent.right
  47.  
  48.                     s.color = x.parent.color
  49.                     x.parent.color = 0
  50.                     s.right.color = 0
  51.                     self.left_rotate(x.parent)
  52.                     x = self.root
  53.             else:
  54.                 s = x.parent.left
  55.                 if s.color == 1:
  56.                     s.color = 0
  57.                     x.parent.color = 1
  58.                     self.right_rotate(x.parent)
  59.                     s = x.parent.left
  60.  
  61.                 if s.right.color == 0 and s.right.color == 0:
  62.                     s.color = 1
  63.                     x = x.parent
  64.                 else:
  65.                     if s.left.color == 0:
  66.                         s.right.color = 0
  67.                         s.color = 1
  68.                         self.left_rotate(s)
  69.                         s = x.parent.left
  70.  
  71.                     s.color = x.parent.color
  72.                     x.parent.color = 0
  73.                     s.left.color = 0
  74.                     self.right_rotate(x.parent)
  75.                     x = self.root
  76.         x.color = 0
  77.  
  78.     def __rb_transplant(self, u, v):
  79.         if u.parent == None:
  80.             self.root = v
  81.         elif u == u.parent.left:
  82.             u.parent.left = v
  83.         else:
  84.             u.parent.right = v
  85.         v.parent = u.parent
  86.  
  87.     def delete_node_helper(self, node, key):
  88.         z = self.TNULL
  89.         while node != self.TNULL:
  90.             if node.item == key:
  91.                 z = node
  92.  
  93.             if node.item <= key:
  94.                 node = node.right
  95.             else:
  96.                 node = node.left
  97.  
  98.         if z == self.TNULL:
  99.             return
  100.  
  101.         y = z
  102.         y_original_color = y.color
  103.         if z.left == self.TNULL:
  104.             x = z.right
  105.             self.__rb_transplant(z, z.right)
  106.         elif (z.right == self.TNULL):
  107.             x = z.left
  108.             self.__rb_transplant(z, z.left)
  109.         else:
  110.             y = self.minimum(z.right)
  111.             y_original_color = y.color
  112.             x = y.right
  113.             if y.parent == z:
  114.                 x.parent = y
  115.             else:
  116.                 self.__rb_transplant(y, y.right)
  117.                 y.right = z.right
  118.                 y.right.parent = y
  119.  
  120.             self.__rb_transplant(z, y)
  121.             y.left = z.left
  122.             y.left.parent = y
  123.             y.color = z.color
  124.         if y_original_color == 0:
  125.             self.delete_fix(x)
  126.  
  127.     def fix_insert(self, k):
  128.         while k.parent.color == 1:
  129.             if k.parent == k.parent.parent.right:
  130.                 u = k.parent.parent.left
  131.                 if u.color == 1:
  132.                     u.color = 0
  133.                     k.parent.color = 0
  134.                     k.parent.parent.color = 1
  135.                     k = k.parent.parent
  136.                 else:
  137.                     if k == k.parent.left:
  138.                         k = k.parent
  139.                         self.right_rotate(k)
  140.                     k.parent.color = 0
  141.                     k.parent.parent.color = 1
  142.                     self.left_rotate(k.parent.parent)
  143.             else:
  144.                 u = k.parent.parent.right
  145.  
  146.                 if u.color == 1:
  147.                     u.color = 0
  148.                     k.parent.color = 0
  149.                     k.parent.parent.color = 1
  150.                     k = k.parent.parent
  151.                 else:
  152.                     if k == k.parent.right:
  153.                         k = k.parent
  154.                         self.left_rotate(k)
  155.                     k.parent.color = 0
  156.                     k.parent.parent.color = 1
  157.                     self.right_rotate(k.parent.parent)
  158.             if k == self.root:
  159.                 break
  160.         self.root.color = 0
  161.  
  162.     def __print_helper(self, node, indent, last):
  163.         if node != self.TNULL:
  164.             sys.stdout.write(indent)
  165.             if last:
  166.                 sys.stdout.write("R----")
  167.                 indent += "     "
  168.             else:
  169.                 sys.stdout.write("L----")
  170.                 indent += "|    "
  171.  
  172.             s_color = "RED" if node.color == 1 else "BLACK"
  173.             print(str(node.item) + "(" + s_color + ")")
  174.             self.__print_helper(node.left, indent, False)
  175.             self.__print_helper(node.right, indent, True)
  176.  
  177.     def preorder(self):
  178.         self.pre_order_helper(self.root)
  179.  
  180.     def minimum(self, node):
  181.         while node.left != self.TNULL:
  182.             node = node.left
  183.         return node
  184.  
  185.     def maximum(self, node):
  186.         while node.right != self.TNULL:
  187.             node = node.right
  188.         return node
  189.  
  190.     def successor(self, x):
  191.         if x.right != self.TNULL:
  192.             return self.minimum(x.right)
  193.  
  194.         y = x.parent
  195.         while y != self.TNULL and x == y.right:
  196.             x = y
  197.             y = y.parent
  198.         return y
  199.  
  200.     def predecessor(self,  x):
  201.         if (x.left != self.TNULL):
  202.             return self.maximum(x.left)
  203.  
  204.         y = x.parent
  205.         while y != self.TNULL and x == y.left:
  206.             x = y
  207.             y = y.parent
  208.  
  209.         return y
  210.  
  211.     def left_rotate(self, x):
  212.         y = x.right
  213.         x.right = y.left
  214.         if y.left != self.TNULL:
  215.             y.left.parent = x
  216.  
  217.         y.parent = x.parent
  218.         if x.parent == None:
  219.             self.root = y
  220.         elif x == x.parent.left:
  221.             x.parent.left = y
  222.         else:
  223.             x.parent.right = y
  224.         y.left = x
  225.         x.parent = y
  226.  
  227.     def right_rotate(self, x):
  228.         y = x.left
  229.         x.left = y.right
  230.         if y.right != self.TNULL:
  231.             y.right.parent = x
  232.  
  233.         y.parent = x.parent
  234.         if x.parent == None:
  235.             self.root = y
  236.         elif x == x.parent.right:
  237.             x.parent.right = y
  238.         else:
  239.             x.parent.left = y
  240.         y.right = x
  241.         x.parent = y
  242.  
  243.     def insert(self, key):
  244.         node = Node(key)
  245.         node.parent = None
  246.         node.item = key
  247.         node.left = self.TNULL
  248.         node.right = self.TNULL
  249.         node.color = 1
  250.  
  251.         y = None
  252.         x = self.root
  253.  
  254.         while x != self.TNULL:
  255.             y = x
  256.             if node.item < x.item:
  257.                 x = x.left
  258.             else:
  259.                 x = x.right
  260.         if y is not None and y.item.key == node.item.key:
  261.             y.item.index.add(node.item.key)
  262.             return
  263.         node.parent = y
  264.         if y == None:
  265.             self.root = node
  266.         elif node.item < y.item:
  267.             y.left = node
  268.         else:
  269.             y.right = node
  270.  
  271.         if node.parent == None:
  272.             node.color = 0
  273.             return
  274.  
  275.         if node.parent.parent == None:
  276.             return
  277.  
  278.         self.fix_insert(node)
  279.  
  280.     def get_root(self):
  281.         return self.root
  282.  
  283.     def delete_node(self, item):
  284.         self.delete_node_helper(self.root, item)
  285.  
  286.     def print_tree(self):
  287.         self.__print_helper(self.root, "", True)
  288.  
  289.  
  290. class ServerStore:
  291.     index: Set[int]
  292.     key: int
  293.  
  294.     def __init__(self, key, ind):
  295.         self.key, self.index = key, set()
  296.         self.index.add(ind)
  297.  
  298.     def __lt__(self, other):
  299.         return self.key < other.key
  300.  
  301.     def __gt__(self, other):
  302.         return self.key > other.key
  303.  
  304.     def __le__(self, other):
  305.         return self.key < other.key
  306.  
  307.     def __ge__(self, other):
  308.         return self.key > other.key
  309.  
  310.     def __eq__(self, other):
  311.         return len(self.index & other.index) > 0
  312.  
  313.     def __str__(self):
  314.         return f"{self.index} {self.key}"
  315.  
  316.  
  317. n, m, q = map(int, input().split())
  318. bt = RedBlackTree()
  319. offs = [set() for _ in range(n)]
  320. off_count = [0] * n
  321. r = [0] * n
  322.  
  323. for i in range(n):
  324.     bt.insert(ServerStore(0, i))
  325.  
  326. for _ in range(q):
  327.     s = input()
  328.     if s.startswith("DISABLE"):
  329.         i, j = map(lambda x: int(x) - 1, s.split()[-2:])
  330.         if j not in offs[i]:
  331.             bt.delete_node_helper(bt.get_root(), ServerStore(r[i] * (m - off_count[i]), i))
  332.             offs[i].add(j)
  333.             off_count[i] += 1
  334.             bt.insert(ServerStore(r[i] * (m - off_count[i]), i))
  335.     elif s.startswith("RESET"):
  336.         i = int(s.split()[-1]) - 1
  337.         bt.delete_node_helper(bt.get_root(), ServerStore(r[i] * (m - off_count[i]), i))
  338.         offs[i] = set()
  339.         off_count[i] = 0
  340.         r[i] += 1
  341.         bt.insert(ServerStore(r[i] * (m - off_count[i]), i))
  342.     elif s.startswith("GETMAX"):
  343.         print(max(bt.maximum(bt.get_root()).item.index) + 1)
  344.     else:
  345.         print(min(bt.minimum(bt.get_root()).item.index) + 1)
  346.  
Advertisement
Add Comment
Please, Sign In to add comment