polysh/rb_tree.py

   1 #!/usr/bin/env python
   2 #
   3 # This code adapted from C source from
   4 # Thomas Niemann's Sorting and Searching Algorithms: A Cookbook
   5 #
   6 # From the title page:
   7 #   Permission to reproduce this document, in whole or in part, is
   8 #   given provided the original web site listed below is referenced,
   9 #   and no additional restrictions apply. Source code, when part of
  10 #   a software project, may be used freely without reference to the
  11 #   author.
  12 #
  13 # http://epaperpress.com
  14 #
  15 # Adapted by Chris Gonnerman <chris.gonnerman@newcenturycomputers.net>
  16 #        and Graham Breed
  17 #
  18 # Adapted by Charles Tolman <ct@acm.org>
  19 #        inheritance from object class
  20 #        added RBTreeIter class
  21 #        added lastNode and prevNode routines to RBTree
  22 #        added RBList class and associated tests
  23
  24 # Trimmed by Guillaume Chazarain <guichaz@gmail.com> to keep only the part
  25 # needed by polysh
  26
  27 __version__ = "1.5-polysh"
  28
  29 import string
  30
  31 BLACK = 0
  32 RED = 1
  33
  34 class RBNode(object):
  35
  36     def __init__(self, key = None, value = None, color = RED):
  37         self.left = self.right = self.parent = None
  38         self.color = color
  39         self.key = key
  40         self.value = value
  41         self.nonzero = 1
  42
  43     def __nonzero__(self):
  44         return self.nonzero
  45
  46 class RBTree(object):
  47
  48     def __init__(self, cmpfn=cmp):
  49         self.sentinel = RBNode()
  50         self.sentinel.left = self.sentinel.right = self.sentinel
  51         self.sentinel.color = BLACK
  52         self.sentinel.nonzero = 0
  53         self.root = self.sentinel
  54         self.count = 0
  55         # changing the comparison function for an existing tree is dangerous!
  56         self.__cmp = cmpfn
  57
  58     def __len__(self):
  59         return self.count
  60
  61     def rotateLeft(self, x):
  62
  63         y = x.right
  64
  65         # establish x.right link
  66         x.right = y.left
  67         if y.left != self.sentinel:
  68             y.left.parent = x
  69
  70         # establish y.parent link
  71         if y != self.sentinel:
  72             y.parent = x.parent
  73         if x.parent:
  74             if x == x.parent.left:
  75                 x.parent.left = y
  76             else:
  77                 x.parent.right = y
  78         else:
  79             self.root = y
  80
  81         # link x and y
  82         y.left = x
  83         if x != self.sentinel:
  84             x.parent = y
  85
  86     def rotateRight(self, x):
  87
  88         #***************************
  89         #  rotate node x to right
  90         #***************************
  91
  92         y = x.left
  93
  94         # establish x.left link
  95         x.left = y.right
  96         if y.right != self.sentinel:
  97             y.right.parent = x
  98
  99         # establish y.parent link
 100         if y != self.sentinel:
 101             y.parent = x.parent
 102         if x.parent:
 103             if x == x.parent.right:
 104                 x.parent.right = y
 105             else:
 106                 x.parent.left = y
 107         else:
 108             self.root = y
 109
 110         # link x and y
 111         y.right = x
 112         if x != self.sentinel:
 113             x.parent = y
 114
 115     def insertFixup(self, x):
 116         #************************************
 117         #  maintain Red-Black tree balance  *
 118         #  after inserting node x           *
 119         #************************************
 120
 121         # check Red-Black properties
 122
 123         while x != self.root and x.parent.color == RED:
 124
 125             # we have a violation
 126
 127             if x.parent == x.parent.parent.left:
 128
 129                 y = x.parent.parent.right
 130
 131                 if y.color == RED:
 132                     # uncle is RED
 133                     x.parent.color = BLACK
 134                     y.color = BLACK
 135                     x.parent.parent.color = RED
 136                     x = x.parent.parent
 137
 138                 else:
 139                     # uncle is BLACK
 140                     if x == x.parent.right:
 141                         # make x a left child
 142                         x = x.parent
 143                         self.rotateLeft(x)
 144
 145                     # recolor and rotate
 146                     x.parent.color = BLACK
 147                     x.parent.parent.color = RED
 148                     self.rotateRight(x.parent.parent)
 149             else:
 150
 151                 # mirror image of above code
 152
 153                 y = x.parent.parent.left
 154
 155                 if y.color == RED:
 156                     # uncle is RED
 157                     x.parent.color = BLACK
 158                     y.color = BLACK
 159                     x.parent.parent.color = RED
 160                     x = x.parent.parent
 161
 162                 else:
 163                     # uncle is BLACK
 164                     if x == x.parent.left:
 165                         x = x.parent
 166                         self.rotateRight(x)
 167
 168                     x.parent.color = BLACK
 169                     x.parent.parent.color = RED
 170                     self.rotateLeft(x.parent.parent)
 171
 172         self.root.color = BLACK
 173
 174     def insertNode(self, key, value):
 175         #**********************************************
 176         #  allocate node for data and insert in tree  *
 177         #**********************************************
 178
 179         # we aren't interested in the value, we just
 180         # want the TypeError raised if appropriate
 181         hash(key)
 182
 183         # find where node belongs
 184         current = self.root
 185         parent = None
 186         while current != self.sentinel:
 187             # GJB added comparison function feature
 188             # slightly improved by JCG: don't assume that ==
 189             # is the same as self.__cmp(..) == 0
 190             rc = self.__cmp(key, current.key)
 191             if rc == 0:
 192                 return current
 193             parent = current
 194             if rc < 0:
 195                 current = current.left
 196             else:
 197                 current = current.right
 198
 199         # setup new node
 200         x = RBNode(key, value)
 201         x.left = x.right = self.sentinel
 202         x.parent = parent
 203
 204         self.count = self.count + 1
 205
 206         # insert node in tree
 207         if parent:
 208             if self.__cmp(key, parent.key) < 0:
 209                 parent.left = x
 210             else:
 211                 parent.right = x
 212         else:
 213             self.root = x
 214
 215         self.insertFixup(x)
 216         return x
 217
 218     def deleteFixup(self, x):
 219         #************************************
 220         #  maintain Red-Black tree balance  *
 221         #  after deleting node x            *
 222         #************************************
 223
 224         while x != self.root and x.color == BLACK:
 225             if x == x.parent.left:
 226                 w = x.parent.right
 227                 if w.color == RED:
 228                     w.color = BLACK
 229                     x.parent.color = RED
 230                     self.rotateLeft(x.parent)
 231                     w = x.parent.right
 232
 233                 if w.left.color == BLACK and w.right.color == BLACK:
 234                     w.color = RED
 235                     x = x.parent
 236                 else:
 237                     if w.right.color == BLACK:
 238                         w.left.color = BLACK
 239                         w.color = RED
 240                         self.rotateRight(w)
 241                         w = x.parent.right
 242
 243                     w.color = x.parent.color
 244                     x.parent.color = BLACK
 245                     w.right.color = BLACK
 246                     self.rotateLeft(x.parent)
 247                     x = self.root
 248
 249             else:
 250                 w = x.parent.left
 251                 if w.color == RED:
 252                     w.color = BLACK
 253                     x.parent.color = RED
 254                     self.rotateRight(x.parent)
 255                     w = x.parent.left
 256
 257                 if w.right.color == BLACK and w.left.color == BLACK:
 258                     w.color = RED
 259                     x = x.parent
 260                 else:
 261                     if w.left.color == BLACK:
 262                         w.right.color = BLACK
 263                         w.color = RED
 264                         self.rotateLeft(w)
 265                         w = x.parent.left
 266
 267                     w.color = x.parent.color
 268                     x.parent.color = BLACK
 269                     w.left.color = BLACK
 270                     self.rotateRight(x.parent)
 271                     x = self.root
 272
 273         x.color = BLACK
 274
 275     def deleteNode(self, z):
 276         #****************************
 277         #  delete node z from tree  *
 278         #****************************
 279
 280         if not z or z == self.sentinel:
 281             return
 282
 283         if z.left == self.sentinel or z.right == self.sentinel:
 284             # y has a self.sentinel node as a child
 285             y = z
 286         else:
 287             # find tree successor with a self.sentinel node as a child
 288             y = z.right
 289             while y.left != self.sentinel:
 290                 y = y.left
 291
 292         # x is y's only child
 293         if y.left != self.sentinel:
 294             x = y.left
 295         else:
 296             x = y.right
 297
 298         # remove y from the parent chain
 299         x.parent = y.parent
 300         if y.parent:
 301             if y == y.parent.left:
 302                 y.parent.left = x
 303             else:
 304                 y.parent.right = x
 305         else:
 306             self.root = x
 307
 308         if y != z:
 309             z.key = y.key
 310             z.value = y.value
 311
 312         if y.color == BLACK:
 313             self.deleteFixup(x)
 314
 315         del y
 316         self.count = self.count - 1
 317
 318     def findNode(self, key):
 319         #******************************
 320         #  find node containing data
 321         #******************************
 322
 323         # we aren't interested in the value, we just
 324         # want the TypeError raised if appropriate
 325         hash(key)
 326
 327         current = self.root
 328
 329         while current != self.sentinel:
 330             # GJB added comparison function feature
 331             # slightly improved by JCG: don't assume that ==
 332             # is the same as self.__cmp(..) == 0
 333             rc = self.__cmp(key, current.key)
 334             if rc == 0:
 335                 return current
 336             else:
 337                 if rc < 0:
 338                     current = current.left
 339                 else:
 340                     current = current.right
 341
 342         return None
 343
 344     def firstNode(self):
 345         cur = self.root
 346         while cur.left:
 347             cur = cur.left
 348         return cur
 349
 350     def lastNode(self):
 351         cur = self.root
 352         while cur.right:
 353             cur = cur.right
 354         return cur
 355