flat-tree/flat_tree/iterator.py
2019-07-01 18:09:49 +02:00

112 lines
2.9 KiB
Python

"""Stateful iterator for flat trees."""
__all__ = ['FlatTreeIterator']
import attr
from flat_tree.accessor import FlatTreeAccessor
@attr.s(auto_attribs=True)
class FlatTreeIterator:
"""Stateful iterator for flat trees."""
index: int = 0
offset: int = 0
factor: int = 0
accessor: FlatTreeAccessor = FlatTreeAccessor()
def __attrs_post_init__(self):
self.seek(self.index)
def next(self) -> int:
"""The next index in the tree."""
self.offset += 1
self.index += self.factor
return self.index
def prev(self) -> int:
"""The previous index in the tree."""
if not self.offset:
return self.index
self.offset -= 1
self.index = self.factor
return self.index
def seek(self, index: int) -> None:
"""Move iterator to the given index.
:param index: The index to move to
"""
self.index = index
if self.index & 1:
self.offset = self.accessor.offset(index)
self.factor = 2 ** (self.accessor.depth(index) + 1)
else:
self.offset = int(index / 2)
self.factor = 2
def parent(self) -> int:
"""Move iterator to the parent index."""
if self.offset & 1:
self.index -= int(self.factor / 2)
self.offset = int((self.offset - 1) / 2)
else:
self.index += int(self.factor / 2)
self.offset = int(self.offset / 2)
self.factor *= 2
return self.index
def left_child(self) -> int:
"""Move iterator to the left child."""
if self.factor == 2:
return self.index
self.factor = int(self.factor / 2)
self.index -= int(self.factor / 2)
self.offset *= 2
return self.index
def right_child(self) -> int:
"""Move iterator to the right child."""
if self.factor == 2:
return self.index
self.factor = int(self.factor / 2)
self.index += int(self.factor / 2)
self.offset = (2 * self.offset) + 1
return self.index
def left_span(self) -> int:
"""Move iterator to the left span."""
self.index = int(self.index - (self.factor / 2)) + 1
self.offset = int(self.index / 2)
self.factor = 2
return self.index
def right_span(self) -> int:
"""Move iterator to the right span."""
self.index = int((self.index - self.factor) / 2) - 1
self.offset = int(self.index / 2)
self.factor = 2
return self.index
def sibling(self) -> int:
"""Move iterator to the sibling."""
return self.next() if self.is_left() else self.prev()
def is_left(self) -> bool:
"""Is this index a left sibling?"""
return not self.offset & 1
def is_right(self) -> bool:
"""Is this index a right sibling?"""
return not self.is_left()