Skip to content

Commit

Permalink
Add get_index utility function
Browse files Browse the repository at this point in the history
  • Loading branch information
joowani committed Feb 13, 2022
1 parent 122ecb6 commit a4bd84c
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 2 deletions.
4 changes: 3 additions & 1 deletion .github/workflows/build.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@ jobs:
run: python -m pip install --upgrade pip setuptools wheel
- name: Install package
run: pip install .[dev]
- name: Check black version
run: black --version
- name: Run black
run: black --check .
run: black -v --check .
- name: Run flake8
run: flake8 .
- name: Run isort
Expand Down
70 changes: 70 additions & 0 deletions binarytree/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"heap",
"build",
"build2",
"get_index",
"get_parent",
"__version__",
]
Expand Down Expand Up @@ -2026,6 +2027,75 @@ def _get_tree_properties(root: Node) -> NodeProperties:
)


def get_index(root: Node, descendent: Node) -> int:
"""Return the level-order_ index given the root and a possible descendent.
.. _level-order:
https://en.wikipedia.org/wiki/Tree_traversal#Breadth-first_search
:return: Level-order index of the descendent relative to the root node.
:rtype: int
:raise binarytree.exceptions.NodeTypeError: If root or descendent is
not an instance of :class:`binarytree.Node`.
:raise binarytree.exceptions.NodeReferenceError: If given a node that is
not a root/descendent.
**Example**:
.. doctest::
>>> from binarytree import Node, get_index
>>>
>>> root = Node(1)
>>> root.left = Node(2)
>>> root.right = Node(3)
>>> root.left.right = Node(4)
>>>
>>> get_index(root, root.left)
1
>>> get_index(root, root.right)
2
>>> get_index(root, root.left.right)
4
>>> get_index(root.left, root.right)
Traceback (most recent call last):
...
NodeReferenceError: given nodes are not in the same tree
"""
if root is None:
raise NodeTypeError("root must be a Node instance")

if descendent is None:
raise NodeTypeError("descendent must be a Node instance")

current_nodes: List[Optional[Node]] = [root]
current_index = 0
has_more_nodes = True

while has_more_nodes:
has_more_nodes = False
next_nodes: List[Optional[Node]] = []

for node in current_nodes:
if node is not None and node is descendent:
return current_index

if node is None:
next_nodes.append(None)
next_nodes.append(None)
else:
next_nodes.append(node.left)
next_nodes.append(node.right)
if node.left is not None or node.right is not None:
has_more_nodes = True

current_index += 1

current_nodes = next_nodes

raise NodeReferenceError("given nodes are not in the same tree")


def get_parent(root: Node, child: Node) -> Optional[Node]:
"""Search the binary tree and return the parent of given child.
Expand Down
5 changes: 5 additions & 0 deletions docs/specs.rst
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,11 @@ Function: binarytree.heap

.. autofunction:: binarytree.heap

Function: binarytree.get_index
===============================

.. autofunction:: binarytree.get_index

Function: binarytree.get_parent
===============================

Expand Down
30 changes: 29 additions & 1 deletion tests/test_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

import pytest

from binarytree import Node, bst, build, build2, get_parent, heap, tree
from binarytree import Node, bst, build, build2, get_index, get_parent, heap, tree
from binarytree.exceptions import (
NodeIndexError,
NodeModifyError,
Expand Down Expand Up @@ -1059,6 +1059,34 @@ def test_heap_float_values():
assert root.size == root_copy.size


def test_get_index():
root = Node(0)
root.left = Node(1)
root.right = Node(2)
root.left.left = Node(3)
root.right.right = Node(4)

assert get_index(root, root) == 0
assert get_index(root, root.left) == 1
assert get_index(root, root.right) == 2
assert get_index(root, root.left.left) == 3
assert get_index(root, root.right.right) == 6

with pytest.raises(NodeReferenceError) as err:
get_index(root.left, root.right)
assert str(err.value) == "given nodes are not in the same tree"

with pytest.raises(NodeTypeError) as err:
# noinspection PyTypeChecker
get_index(root, None)
assert str(err.value) == "descendent must be a Node instance"

with pytest.raises(NodeTypeError) as err:
# noinspection PyTypeChecker
get_index(None, root.left)
assert str(err.value) == "root must be a Node instance"


# noinspection PyTypeChecker
def test_get_parent():
root = Node(0)
Expand Down

0 comments on commit a4bd84c

Please sign in to comment.