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