- Please read over PEP 8 as your naming style is unidiomatic, and would cause programmers to be confused when seeing your code in the wild.
- I would move getting and setting a node by index into their own functions.
if self.value >= 0 only hinders your code, the value -1 is perfectly valid in binary trees. It also means that you're limiting value to types that can be compared to integers, meaning you can't enter strings.
- Using recursion in
addNode is a good idea, but I find the non-recursive alternate to be easier to understand.
- Your code doesn't care if I enter
addNode([0, 2], 1). This seems fishy as if this were a Trinary Tree that would mean something completely different. I recommend raising a ValueError in this case.
- Your
printTree is pretty good.
- I'd change the
level argument in printTree to be a keyword argument only. This is because then it's explicit that it's changing the level, and it's something normal code shouldn't accidentally change.
- Personally I think
if self.left is not None is better here than if self.left, but as Node can't be falsy it doesn't matter too much.
- I've added some basic docstrings, defined in PEP 257.
- I have also added some
typing to your code.
- By using
typing and mypy I found that addNode can error when the specified path doesn't exist. It would be better in these cases to raise an error telling the user why.
- You can use a
if __name__ == '__main__': guard to prevent code running on import. This means if you later import this code in another module then the testing code won't run.
from __future__ import annotations
from typing import Any, Optional, Sequence
class Node:
"""Test binary tree."""
def __init__(
self,
value: Any,
right: Optional[Node] = None,
left: Optional[Node] = None,
):
"""Initialize binary tree, with specified children nodes."""
self.value = value
self.right = right
self.left = left
def _get_index(self, index: int) -> Optional[Node]:
"""Get node via integer index."""
if index == 0:
return self.left
elif index == 1:
return self.right
else:
raise ValueError(f'Invalid index {index}')
def _set_index(self, index: int, value: Node) -> None:
"""Set node via integer index."""
if index == 0:
self.left = value
elif index == 1:
self.right = value
else:
raise ValueError(f'Invalid index {index}')
def add_node(self, parents: Sequence[int], value: Node) -> None:
"""Add the provided node to the tree."""
node: Optional[Node] = self
for index in parents[:-1]:
if node is not None:
node = node._get_index(index)
if node is None:
raise ValueError("Parent node doesn't exist.")
node._set_index(parents[-1], value)
def add_value(self, parents: Sequence[int], value: Any) -> None:
"""Add the provided value to the tree."""
self.add_node(parents, Node(value))
def print_tree(self, *, level: int = 0) -> None:
"""Print the tree."""
print(' ' * level + str(self.value))
for child in (self.left, self.right):
if child is not None:
child.print_tree(level=level + 1)
if __name__ == '__main__':
tree = Node(0)
tree.add_value([0], 1)
tree.add_value([1], 1)
r"""
Graph looks like
0
/ \
1 1
"""
tree.add_value([0, 0], 3)
tree.add_value([0, 1], 8)
tree.add_value([1, 0], 4)
tree.add_value([1, 1], 7)
r"""
Graph looks like
0
/ \
1 1
/ \ / \
3 8 4 7
"""
tree.print_tree()