fork download
copy 
  1. #!/usr/bin/python
  2. class rbnode(object):
  3.  
  4.     def __init__(self, key):
  5.         self._key = key
  6.         self._red = False
  7.         self._left = None
  8.         self._right = None
  9.         self._p = None
  10.  
  11.  
  12.  
  13. class rbtree(object):
  14.  
  15.  
  16.     def __init__(self, create_node=rbnode):
  17.  
  18.         self._nil = create_node(key=None)
  19.  
  20.         self._root = self.nil
  21.  
  22.         self._create_node = create_node
  23.  
  24.  
  25.  
  26.     def insert_key(self, key):
  27.         self.insert_node(self._create_node(key=key))
  28.  
  29.  
  30.     def insert_node(self, z):
  31.         y = self.nil
  32.         x = self.root
  33.         while x != self.nil:
  34.             y = x
  35.             if z.key < x.key:
  36.                 x = x.left
  37.             else:
  38.                 x = x.right
  39.         z._p = y
  40.         if y == self.nil:
  41.             self._root = z
  42.         elif z.key < y.key:
  43.             y._left = z
  44.         else:
  45.             y._right = z
  46.         z._left = self.nil
  47.         z._right = self.nil
  48.         z._red = True
  49.         self._insert_fixup(z)
  50.  
  51.  
  52.     def _insert_fixup(self, z):
  53.         while z.p.red:
  54.             if z.p == z.p.p.left:
  55.                 y = z.p.p.right
  56.                 if y.red:
  57.                     z.p._red = False
  58.                     y._red = False
  59.                     z.p.p._red = True
  60.                     z = z.p.p
  61.                 else:
  62.                     if z == z.p.right:
  63.                         z = z.p
  64.                         self._left_rotate(z)
  65.                     z.p._red = False
  66.                     z.p.p._red = True
  67.                     self._right_rotate(z.p.p)
  68.             else:
  69.                 y = z.p.p.left
  70.                 if y.red:
  71.                     z.p._red = False
  72.                     y._red = False
  73.                     z.p.p._red = True
  74.                     z = z.p.p
  75.                 else:
  76.                     if z == z.p.left:
  77.                         z = z.p
  78.                         self._right_rotate(z)
  79.                     z.p._red = False
  80.                     z.p.p._red = True
  81.                     self._left_rotate(z.p.p)
  82.         self.root._red = False
  83.  
  84.  
  85.     def _left_rotate(self, x):
  86.         y = x.right
  87.         x._right = y.left
  88.         if y.left != self.nil:
  89.             y.left._p = x
  90.         y._p = x.p
  91.         if x.p == self.nil:
  92.             self._root = y
  93.         elif x == x.p.left:
  94.             x.p._left = y
  95.         else:
  96.             x.p._right = y
  97.         y._left = x
  98.         x._p = y
  99.  
  100.  
  101.     def _right_rotate(self, y):
  102.         x = y.left
  103.         y._left = x.right
  104.         if x.right != self.nil:
  105.             x.right._p = y
  106.         x._p = y.p
  107.         if y.p == self.nil:
  108.             self._root = x
  109.         elif y == y.p.right:
  110.             y.p._right = x
  111.         else:
  112.             y.p._left = x
  113.         x._right = y
  114.         y._p = x
  115.  
  116.  
  117.     def check_invariants(self):
  118.  
  119.         def is_red_black_node(node):
  120.             if (node.left and not node.right) or (node.right and not node.left):
  121.                 return 0, False
  122.  
  123.             # check leaves are black
  124.             if not node.left and not node.right and node.red:
  125.                 return 0, False
  126.  
  127.             # if node is red, check children are black
  128.             if node.red and node.left and node.right:
  129.                 if node.left.red or node.right.red:
  130.                     return 0, False
  131.  
  132.             # descend tree and check black counts are balanced
  133.             if node.left and node.right:
  134.  
  135.                 # check children's parents are correct
  136.                 if self.nil != node.left and node != node.left.p:
  137.                     return 0, False
  138.                 if self.nil != node.right and node != node.right.p:
  139.                     return 0, False
  140.  
  141.                 # check children are ok
  142.                 left_counts, left_ok = is_red_black_node(node.left)
  143.                 if not left_ok:
  144.                     return 0, False
  145.                 right_counts, right_ok = is_red_black_node(node.right)
  146.                 if not right_ok:
  147.                     return 0, False
  148.  
  149.                 # check children's counts are ok
  150.                 if left_counts != right_counts:
  151.                     return 0, False
  152.                 return left_counts, True
  153.             else:
  154.                 return 0, True
  155.  
  156.         num_black, is_ok = is_red_black_node(self.root)
  157.         return is_ok and not self.root._red
  158.  
  159.     def node_color(node):
  160.         if node.red:
  161.             return "red"
  162.         else:
  163.             return "black"
  164.  
  165.     def visit_node(node):
  166.  
  167.         if node.left:
  168.             if node.left != t.nil or show_nil:
  169.                 visit_node(node.left)
  170.                 print ("%s,%s,%s"%(node.key,node.color,node.p))
  171.         if node.right:
  172.             if node.right != t.nil or show_nil:
  173.                 visit_node(node.right)
  174.                 print ("%s,%s,%s"%(node.key,node.color,node.p))
  175.  
  176.         visit_node(t.root)
  177.  
  178.  
  179.  
  180.  
  181. def test_tree(t, keys):
  182.     assert t.check_invariants()
  183.     for i, key in enumerate(keys):
  184.         for key2 in keys[:i]:
  185.             assert t.nil != t.search(key2)
  186.         for key2 in keys[i:]:
  187.             assert (t.nil == t.search(key2)) ^ (key2 in keys[:i])
  188.         t.insert_key(key)
  189.         assert t.check_invariants()
  190.  
  191.  
  192.     # test the rbtree
  193. n = int(input())
  194. keys = [int(x) for x in input().split()]   
  195. size=n
  196. t = rbtree()
  197. for x in keys :
  198.     test_tree(t, x)
  199.  
  200. # your code goes here
Runtime error #stdin #stdout #stderr 0s 23304KB
comments (?)
stdin
copy 
5 1 2 3 4 5
stdout
copy 
Standard output is empty
stderr
copy 
Traceback (most recent call last):
  File "prog.py", line 193, in <module>
  File "<string>", line 1
    5 1 2 3 4 5
      ^
SyntaxError: invalid syntax
https://ideone.com/SS87Iq
language:
Python (cpython 2.7.16)
created:
6 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!