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 9from operator import attrgetter, lt, gt 10 11 12class 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 221