third_party/bintrees/bintrees/qrbtree.pyx

   1 #!/usr/bin/env python
   2 #coding:utf-8
   3 # Author:  mozman
   4 # Purpose: cython unbalanced binary tree module
   5 # Created: 28.04.2010
   6 # Copyright (c) 2010-2013 by Manfred Moitzi
   7 # License: MIT License
   8
   9 __all__ = ['cRBTree']
  10
  11 from cwalker import cWalker
  12
  13 from cwalker cimport *
  14 from ctrees cimport *
  15
  16 cdef class cRBTree:
  17     cdef node_t *_root
  18     cdef int _count
  19
  20     def __cinit__(self, items=None):
  21         self._root = NULL
  22         self._count = 0
  23         if items:
  24             self.update(items)
  25
  26     def __dealloc__(self):
  27         ct_delete_tree(self._root)
  28
  29     @property
  30     def count(self):
  31         return self._count
  32
  33     def __getstate__(self):
  34         return dict(self.items())
  35
  36     def __setstate__(self, state):
  37         self.update(state)
  38
  39     def clear(self):
  40         ct_delete_tree(self._root)
  41         self._count = 0
  42         self._root = NULL
  43
  44     def get_value(self, key):
  45         result = <object> ct_get_item(self._root, key)
  46         if result is None:
  47             raise KeyError(key)
  48         else:
  49             return result[1]
  50
  51     def get_walker(self):
  52         cdef cWalker walker
  53         walker = cWalker()
  54         walker.set_tree(self._root)
  55         return walker
  56
  57     def insert(self, key, value):
  58         res = rb_insert(&self._root, key, value)
  59         if res < 0:
  60             raise MemoryError('Can not allocate memory for node structure.')
  61         else:
  62             self._count += res
  63
  64     def remove(self, key):
  65         cdef int result
  66         result =  rb_remove(&self._root, key)
  67         if result == 0:
  68             raise KeyError(str(key))
  69         else:
  70             self._count -= 1
  71
  72     def max_item(self):
  73         """ Get item with max key of tree, raises ValueError if tree is empty. """
  74         cdef node_t *node
  75         node = ct_max_node(self._root)
  76         if node == NULL: # root is None
  77             raise ValueError("Tree is empty")
  78         return (<object>node.key, <object>node.value)
  79
  80     def min_item(self):
  81         """ Get item with min key of tree, raises ValueError if tree is empty. """
  82         cdef node_t *node
  83         node = ct_min_node(self._root)
  84         if node == NULL: # root is None
  85             raise ValueError("Tree is empty")
  86         return (<object>node.key, <object>node.value)