flat-tree implementation coming down the tubes
This commit is contained in:
Luke Murphy
15
flat_tree/__init__.py
Normal file
15
flat_tree/__init__.py
Normal file
@ -0,0 +1,15 @@
|
||||
"""Flat tree module."""
|
||||
|
||||
from flat_tree.accessor import FlatTreeAccessor # noqa
|
||||
from flat_tree.iterator import FlatTreeIterator # noqa
|
||||
|
||||
try:
|
||||
import pkg_resources
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
__version__ = pkg_resources.get_distribution('flat_tree').version
|
||||
except Exception:
|
||||
__version__ = 'unknown'
|
||||
205
flat_tree/accessor.py
Normal file
205
flat_tree/accessor.py
Normal file
@ -0,0 +1,205 @@
|
||||
"""An accessor for navigating flat trees."""
|
||||
|
||||
__all__ = ['FlatTreeAccessor']
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
class FlatTreeAccessor:
|
||||
"""A flat tree accessor."""
|
||||
|
||||
def index(self, depth: int, offset: int) -> int:
|
||||
"""The tree index specified by the depth and offset.
|
||||
|
||||
:param depth: The depth of the tree
|
||||
:param offset: The offset from left hand side of the tree
|
||||
"""
|
||||
return ((1 + (2 * offset)) * (2 ** depth)) - 1
|
||||
|
||||
def offset(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The offset of given index from the left hand side of the tree.
|
||||
|
||||
:param index: The tree index
|
||||
:param depth: The depth of the tree
|
||||
"""
|
||||
if not (index & 1):
|
||||
return int(index / 2)
|
||||
|
||||
if depth is None:
|
||||
depth = self.depth(index)
|
||||
|
||||
return int((((index + 1) / (2 ** depth)) - 1) / 2)
|
||||
|
||||
def depth(self, index: int) -> int:
|
||||
"""The depth of the given index in the tree.
|
||||
|
||||
:param index: The tree index
|
||||
"""
|
||||
depth = 0
|
||||
|
||||
index += 1
|
||||
while not (index & 1):
|
||||
depth += 1
|
||||
index = index >> 1
|
||||
|
||||
return depth
|
||||
|
||||
def parent(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The index of the parent relative to the given index.
|
||||
|
||||
:param index: The index relative to the parent
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if depth is None:
|
||||
depth = self.depth(index)
|
||||
|
||||
offset = self.offset(index, depth)
|
||||
|
||||
return self.index(depth + 1, int((offset - (offset & 1)) / 2))
|
||||
|
||||
def sibling(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The index of the sibling relative to the given index.
|
||||
|
||||
:param index: The index relative to the sibling
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if depth is None:
|
||||
depth = self.depth(index)
|
||||
|
||||
offset = self.offset(index, depth)
|
||||
offset = offset - 1 if (offset & 1) else offset + 1
|
||||
|
||||
return self.index(depth, offset)
|
||||
|
||||
def children(self, index: int, depth: Optional[int] = None) -> List[int]:
|
||||
"""All children relative to the given index.
|
||||
|
||||
:param index: The parent index
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if not (index & 1):
|
||||
return []
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
offset = self.offset(index, depth) * 2
|
||||
|
||||
return [
|
||||
self.index((depth - 1), offset),
|
||||
self.index((depth - 1), (offset + 1)),
|
||||
]
|
||||
|
||||
def spans(self, index: int, depth: Optional[int] = None) -> List[int]:
|
||||
"""The span of the tree.
|
||||
|
||||
:param index: The index of the root
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if not (index & 1):
|
||||
return [index, index]
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
offset = self.offset(index, depth)
|
||||
width = 2 ** (depth + 1)
|
||||
|
||||
return [(offset * width), ((offset + 1) * width) - 2]
|
||||
|
||||
def left_span(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The leftmost span of the tree.
|
||||
|
||||
:param index: The index of the tree root
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if not (index & 1):
|
||||
return index
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
return self.offset(index, depth) * (2 ** (depth + 1))
|
||||
|
||||
def right_span(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The rightmost span of the tree.
|
||||
|
||||
:param index: The index of the tree root
|
||||
:param depth: The depth of the index
|
||||
"""
|
||||
if not (index & 1):
|
||||
return index
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
return (self.offset(index, depth) + 1) * (2 ** (depth + 1)) - 2
|
||||
|
||||
def count(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The number of nodes a tree contains.
|
||||
|
||||
:param index: The index of the root of the tree
|
||||
:param depth: The depth of the root of the tree
|
||||
"""
|
||||
if not (index & 1):
|
||||
return 1
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
return (2 ** (depth + 1)) - 1
|
||||
|
||||
def full_roots(self, index: int) -> List[int]:
|
||||
"""All full roots within the tree.
|
||||
|
||||
:param index: The index of the root of the tree
|
||||
"""
|
||||
if index & 1:
|
||||
message = 'Roots only available for tree depth 0'
|
||||
raise ValueError(message)
|
||||
|
||||
roots: List[int] = []
|
||||
|
||||
index = int(index / 2)
|
||||
offset, factor = 0, 1
|
||||
|
||||
while True:
|
||||
if not index:
|
||||
return roots
|
||||
|
||||
while (factor * 2) <= index:
|
||||
factor = factor * 2
|
||||
|
||||
roots.append((offset + factor) - 1)
|
||||
|
||||
offset = (offset + 2) * factor
|
||||
index = index - factor
|
||||
factor = 1
|
||||
|
||||
def left_child(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The left child of the given index.
|
||||
|
||||
:param index: The index of the tree
|
||||
:param depth: The depth of the tree
|
||||
"""
|
||||
if not (index & 1):
|
||||
return -1
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
return self.index((depth - 1), (self.offset(index, depth) * 2))
|
||||
|
||||
def right_child(self, index: int, depth: Optional[int] = None) -> int:
|
||||
"""The right child of the given index.
|
||||
|
||||
:param index: The index of the tree
|
||||
:param depth: The depth of the tree
|
||||
"""
|
||||
if not (index & 1):
|
||||
return -1
|
||||
|
||||
if not depth:
|
||||
depth = self.depth(index)
|
||||
|
||||
return self.index((depth - 1), (1 + (self.offset(index, depth) * 2)))
|
||||
111
flat_tree/iterator.py
Normal file
111
flat_tree/iterator.py
Normal file
@ -0,0 +1,111 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user