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