third_party/bintrees/bintrees/walker.py

   1 #!/usr/bin/env python
   2 #coding:utf-8
   3 # Author:  mozman
   4 # Purpose: tree walker
   5 # Created: 07.05.2010
   6 # Copyright (c) 2010-2013 by Manfred Moitzi
   7 # License: MIT License
   8
   9 from operator import attrgetter, lt, gt
  10
  11
  12 class Walker(object):
  13     __slots__ = ['_node', '_stack', '_tree']
  14
  15     def __init__(self, tree):
  16         self._tree = tree
  17         self._node = tree.root
  18         self._stack = []
  19
  20     def reset(self):
  21         self._stack = []
  22         self._node = self._tree.root
  23
  24     @property
  25     def key(self):
  26         return self._node.key
  27
  28     @property
  29     def value(self):
  30         return self._node.value
  31
  32     @property
  33     def item(self):
  34         return (self._node.key, self._node.value)
  35
  36     @property
  37     def is_valid(self):
  38         return self._node is not None
  39
  40     def goto(self, key):
  41         self._node = self._tree.root
  42         while self._node is not None:
  43             if key == self._node.key:
  44                 return True
  45             elif key < self._node.key:
  46                 self.go_left()
  47             else:
  48                 self.go_right()
  49         return False
  50
  51     def push(self):
  52         self._stack.append(self._node)
  53
  54     def pop(self):
  55         self._node = self._stack.pop()
  56
  57     def stack_is_empty(self):
  58         return len(self._stack) == 0
  59
  60     def goto_leaf(self):
  61         """ get a leaf node """
  62         while self._node is not None:
  63             if self.has_left():
  64                 self.go_left()
  65             elif self.has_right():
  66                 self.go_right()
  67             else:
  68                 return
  69
  70     def has_child(self, direction):
  71         if direction == 0:
  72             return self._node.left is not None
  73         else:
  74             return self._node.right is not None
  75
  76     def down(self, direction):
  77         if direction == 0:
  78             self._node = self._node.left
  79         else:
  80             self._node = self._node.right
  81
  82     def go_left(self):
  83         self._node = self._node.left
  84
  85     def go_right(self):
  86         self._node = self._node.right
  87
  88     def has_left(self):
  89         return self._node.left is not None
  90
  91     def has_right(self):
  92         return self._node.right is not None
  93
  94     def _next_item(self, key, left, right, less_than):
  95         node = self._tree.root
  96         succ = None
  97         while node is not None:
  98             if key == node.key:
  99                 break
 100             elif less_than(key, node.key):
 101                 if (succ is None) or less_than(node.key, succ.key):
 102                     succ = node
 103                 node = left(node)
 104             else:
 105                 node = right(node)
 106
 107         if node is None:  # stay at dead end
 108             raise KeyError(str(key))
 109         # found node of key
 110         if right(node) is not None:
 111             # find smallest node of right subtree
 112             node = right(node)
 113             while left(node) is not None:
 114                 node = left(node)
 115             if succ is None:
 116                 succ = node
 117             elif less_than(node.key, succ.key):
 118                 succ = node
 119         elif succ is None:  # given key is biggest in tree
 120             raise KeyError(str(key))
 121         return (succ.key, succ.value)
 122
 123     def succ_item(self, key):
 124         """ Get successor (k,v) pair of key, raises KeyError if key is max key
 125         or key does not exist.
 126         """
 127         return self._next_item(
 128             key,
 129             left=attrgetter("left"),
 130             right=attrgetter("right"),
 131             less_than=lt,
 132         )
 133
 134     def prev_item(self, key):
 135         """ Get predecessor (k,v) pair of key, raises KeyError if key is min key
 136         or key does not exist.
 137         """
 138         return self._next_item(
 139             key,
 140             left=attrgetter("right"),
 141             right=attrgetter("left"),
 142             less_than=gt,
 143         )
 144
 145     def _iteritems(self, left=attrgetter("left"), right=attrgetter("right")):
 146         """ optimized forward iterator (reduced method calls) """
 147         if self._tree.is_empty():
 148             return
 149         node = self._tree.root
 150         stack = self._stack
 151         go_left = True
 152         while True:
 153             if left(node) is not None and go_left:
 154                 stack.append(node)
 155                 node = left(node)
 156             else:
 157                 yield (node.key, node.value)
 158                 if right(node) is not None:
 159                     node = right(node)
 160                     go_left = True
 161                 else:
 162                     if len(stack) == 0:
 163                         return  # all done
 164                     node = stack.pop()
 165                     go_left = False
 166
 167     def iter_items_forward(self):
 168         for item in self._iteritems(
 169             left=attrgetter("left"),
 170             right=attrgetter("right"),
 171         ):
 172             yield item
 173
 174     def iter_items_backward(self):
 175         for item in self._iteritems(
 176             left=attrgetter("right"),
 177             right=attrgetter("left"),
 178         ):
 179             yield item
 180
 181     def floor_item(self, key):
 182         """ Get the element (k,v) pair associated with the greatest key less
 183         than or equal to the given key, raises KeyError if there is no such key.
 184         """
 185         node = self._tree.root
 186         prev = None
 187         while node is not None:
 188             if key == node.key:
 189                 return node.key, node.value
 190             elif key < node.key:
 191                 node = node.left
 192             else:
 193                 if (prev is None) or (node.key > prev.key):
 194                     prev = node
 195                 node = node.right
 196         # node must be None here
 197         if prev:
 198             return prev.key, prev.value
 199         raise KeyError(str(key))
 200
 201     def ceiling_item(self, key):
 202         """ Get the element (k,v) pair associated with the smallest key greater
 203         than or equal to the given key, raises KeyError if there is no such key.
 204         """
 205         node = self._tree.root
 206         succ = None
 207         while node is not None:
 208             if key == node.key:
 209                 return node.key, node.value
 210             elif key > node.key:
 211                 node = node.right
 212             else:
 213                 if (succ is None) or (node.key < succ.key):
 214                     succ = node
 215                 node = node.left
 216             # node must be None here
 217         if succ:
 218             return succ.key, succ.value
 219         raise KeyError(str(key))
 220