419 lines
14 KiB
Python
419 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, List, Optional, Tuple, TypeVar, Union
|
|
|
|
from bigtree.globals import ASSERTIONS
|
|
from bigtree.node.node import Node
|
|
from bigtree.utils.exceptions import CorruptedTreeError, LoopError, TreeError
|
|
|
|
|
|
class BinaryNode(Node):
|
|
"""
|
|
BinaryNode is an extension of Node, and is able to extend to any Python class for Binary Tree implementation.
|
|
Nodes can have attributes if they are initialized from `BinaryNode`, *dictionary*, or *pandas DataFrame*.
|
|
|
|
BinaryNode can be linked to each other with `children`, `left`, or `right` setter methods.
|
|
If initialized with `children`, it must be length 2, denoting left and right child.
|
|
|
|
Examples:
|
|
>>> from bigtree import BinaryNode, print_tree
|
|
>>> a = BinaryNode(1)
|
|
>>> b = BinaryNode(2)
|
|
>>> c = BinaryNode(3)
|
|
>>> d = BinaryNode(4)
|
|
>>> a.children = [b, c]
|
|
>>> b.right = d
|
|
>>> print_tree(a)
|
|
1
|
|
├── 2
|
|
│ └── 4
|
|
└── 3
|
|
|
|
Directly passing `left`, `right`, or `children` argument.
|
|
|
|
>>> from bigtree import BinaryNode
|
|
>>> d = BinaryNode(4)
|
|
>>> c = BinaryNode(3)
|
|
>>> b = BinaryNode(2, right=d)
|
|
>>> a = BinaryNode(1, children=[b, c])
|
|
|
|
**BinaryNode Creation**
|
|
|
|
Node can be created by instantiating a `BinaryNode` class or by using a *dictionary*.
|
|
If node is created with dictionary, all keys of dictionary will be stored as class attributes.
|
|
|
|
>>> from bigtree import BinaryNode
|
|
>>> a = BinaryNode.from_dict({"name": "1"})
|
|
>>> a
|
|
BinaryNode(name=1, val=1)
|
|
|
|
**BinaryNode Attributes**
|
|
|
|
These are node attributes that have getter and/or setter methods.
|
|
|
|
Get `BinaryNode` configuration
|
|
|
|
1. ``left``: Get left children
|
|
2. ``right``: Get right children
|
|
|
|
----
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: Union[str, int] = "",
|
|
left: Optional[T] = None,
|
|
right: Optional[T] = None,
|
|
parent: Optional[T] = None,
|
|
children: Optional[List[Optional[T]]] = None,
|
|
**kwargs: Any,
|
|
):
|
|
try:
|
|
self.val: Union[str, int] = int(name)
|
|
except ValueError:
|
|
self.val = str(name)
|
|
self.name = str(name)
|
|
self._sep = "/"
|
|
self.__parent: Optional[T] = None
|
|
self.__children: List[Optional[T]] = [None, None]
|
|
if not children:
|
|
children = []
|
|
if len(children):
|
|
if len(children) and len(children) != 2:
|
|
raise ValueError("Children input must have length 2")
|
|
if left and left != children[0]:
|
|
raise ValueError(
|
|
f"Error setting child: Attempting to set both left and children with mismatched values\n"
|
|
f"Check left {left} and children {children}"
|
|
)
|
|
if right and right != children[1]:
|
|
raise ValueError(
|
|
f"Error setting child: Attempting to set both right and children with mismatched values\n"
|
|
f"Check right {right} and children {children}"
|
|
)
|
|
else:
|
|
children = [left, right]
|
|
self.parent = parent
|
|
self.children = children # type: ignore
|
|
if "parents" in kwargs:
|
|
raise AttributeError(
|
|
"Attempting to set `parents` attribute, do you mean `parent`?"
|
|
)
|
|
self.__dict__.update(**kwargs)
|
|
|
|
@property
|
|
def left(self: T) -> T:
|
|
"""Get left children
|
|
|
|
Returns:
|
|
(Self)
|
|
"""
|
|
return self.__children[0]
|
|
|
|
@left.setter
|
|
def left(self: T, left_child: Optional[T]) -> None:
|
|
"""Set left children
|
|
|
|
Args:
|
|
left_child (Optional[Self]): left child
|
|
"""
|
|
self.children = [left_child, self.right] # type: ignore
|
|
|
|
@property
|
|
def right(self: T) -> T:
|
|
"""Get right children
|
|
|
|
Returns:
|
|
(Self)
|
|
"""
|
|
return self.__children[1]
|
|
|
|
@right.setter
|
|
def right(self: T, right_child: Optional[T]) -> None:
|
|
"""Set right children
|
|
|
|
Args:
|
|
right_child (Optional[Self]): right child
|
|
"""
|
|
self.children = [self.left, right_child] # type: ignore
|
|
|
|
@staticmethod
|
|
def __check_parent_type(new_parent: Optional[T]) -> None:
|
|
"""Check parent type
|
|
|
|
Args:
|
|
new_parent (Optional[Self]): parent node
|
|
"""
|
|
if not (isinstance(new_parent, BinaryNode) or new_parent is None):
|
|
raise TypeError(
|
|
f"Expect parent to be BinaryNode type or NoneType, received input type {type(new_parent)}"
|
|
)
|
|
|
|
@property
|
|
def parent(self: T) -> Optional[T]:
|
|
"""Get parent node
|
|
|
|
Returns:
|
|
(Optional[Self])
|
|
"""
|
|
return self.__parent
|
|
|
|
@parent.setter
|
|
def parent(self: T, new_parent: Optional[T]) -> None:
|
|
"""Set parent node
|
|
|
|
Args:
|
|
new_parent (Optional[Self]): parent node
|
|
"""
|
|
if ASSERTIONS:
|
|
self.__check_parent_type(new_parent)
|
|
self._BaseNode__check_parent_loop(new_parent) # type: ignore
|
|
|
|
current_parent = self.parent
|
|
current_child_idx = None
|
|
|
|
# Assign new parent - rollback if error
|
|
self.__pre_assign_parent(new_parent)
|
|
try:
|
|
# Remove self from old parent
|
|
if current_parent is not None:
|
|
if not any(
|
|
child is self for child in current_parent.children
|
|
): # pragma: no cover
|
|
raise CorruptedTreeError(
|
|
"Error setting parent: Node does not exist as children of its parent"
|
|
)
|
|
current_child_idx = current_parent.__children.index(self)
|
|
current_parent.__children[current_child_idx] = None
|
|
|
|
# Assign self to new parent
|
|
self.__parent = new_parent
|
|
if new_parent is not None:
|
|
inserted = False
|
|
for child_idx, child in enumerate(new_parent.__children):
|
|
if not child and not inserted:
|
|
new_parent.__children[child_idx] = self
|
|
inserted = True
|
|
if not inserted:
|
|
raise TreeError(f"Parent {new_parent} already has 2 children")
|
|
|
|
self.__post_assign_parent(new_parent)
|
|
|
|
except Exception as exc_info:
|
|
# Remove self from new parent
|
|
if new_parent is not None and self in new_parent.__children:
|
|
child_idx = new_parent.__children.index(self)
|
|
new_parent.__children[child_idx] = None
|
|
|
|
# Reassign self to old parent
|
|
self.__parent = current_parent
|
|
if current_child_idx is not None:
|
|
current_parent.__children[current_child_idx] = self
|
|
raise TreeError(exc_info)
|
|
|
|
def __pre_assign_parent(self: T, new_parent: Optional[T]) -> None:
|
|
"""Custom method to check before attaching parent
|
|
Can be overridden with `_BinaryNode__pre_assign_parent()`
|
|
|
|
Args:
|
|
new_parent (Optional[Self]): new parent to be added
|
|
"""
|
|
pass
|
|
|
|
def __post_assign_parent(self: T, new_parent: Optional[T]) -> None:
|
|
"""Custom method to check after attaching parent
|
|
Can be overridden with `_BinaryNode__post_assign_parent()`
|
|
|
|
Args:
|
|
new_parent (Optional[Self]): new parent to be added
|
|
"""
|
|
pass
|
|
|
|
def __check_children_type(
|
|
self: T, new_children: List[Optional[T]]
|
|
) -> List[Optional[T]]:
|
|
"""Check child type
|
|
|
|
Args:
|
|
new_children (List[Optional[Self]]): child node
|
|
|
|
Returns:
|
|
(List[Optional[Self]])
|
|
"""
|
|
if not len(new_children):
|
|
new_children = [None, None]
|
|
if len(new_children) != 2:
|
|
raise ValueError("Children input must have length 2")
|
|
return new_children
|
|
|
|
def __check_children_loop(self: T, new_children: List[Optional[T]]) -> None:
|
|
"""Check child loop
|
|
|
|
Args:
|
|
new_children (List[Optional[Self]]): child node
|
|
"""
|
|
seen_children = []
|
|
for new_child in new_children:
|
|
# Check type
|
|
if new_child is not None and not isinstance(new_child, BinaryNode):
|
|
raise TypeError(
|
|
f"Expect children to be BinaryNode type or NoneType, received input type {type(new_child)}"
|
|
)
|
|
|
|
# Check for loop and tree structure
|
|
if new_child is self:
|
|
raise LoopError("Error setting child: Node cannot be child of itself")
|
|
if any(child is new_child for child in self.ancestors):
|
|
raise LoopError(
|
|
"Error setting child: Node cannot be ancestor of itself"
|
|
)
|
|
|
|
# Check for duplicate children
|
|
if new_child is not None:
|
|
if id(new_child) in seen_children:
|
|
raise TreeError(
|
|
"Error setting child: Node cannot be added multiple times as a child"
|
|
)
|
|
else:
|
|
seen_children.append(id(new_child))
|
|
|
|
@property
|
|
def children(self: T) -> Tuple[T, ...]:
|
|
"""Get child nodes
|
|
|
|
Returns:
|
|
(Tuple[Optional[Self]])
|
|
"""
|
|
return tuple(self.__children)
|
|
|
|
@children.setter
|
|
def children(self: T, _new_children: List[Optional[T]]) -> None:
|
|
"""Set child nodes
|
|
|
|
Args:
|
|
_new_children (List[Optional[Self]]): child node
|
|
"""
|
|
self._BaseNode__check_children_type(_new_children) # type: ignore
|
|
new_children = self.__check_children_type(_new_children)
|
|
if ASSERTIONS:
|
|
self.__check_children_loop(new_children)
|
|
|
|
current_new_children = {
|
|
new_child: (
|
|
new_child.parent.__children.index(new_child),
|
|
new_child.parent,
|
|
)
|
|
for new_child in new_children
|
|
if new_child is not None and new_child.parent is not None
|
|
}
|
|
current_new_orphan = [
|
|
new_child
|
|
for new_child in new_children
|
|
if new_child is not None and new_child.parent is None
|
|
]
|
|
current_children = list(self.children)
|
|
|
|
# Assign new children - rollback if error
|
|
self.__pre_assign_children(new_children)
|
|
try:
|
|
# Remove old children from self
|
|
del self.children
|
|
|
|
# Assign new children to self
|
|
self.__children = new_children
|
|
for new_child in new_children:
|
|
if new_child is not None:
|
|
if new_child.parent:
|
|
child_idx = new_child.parent.__children.index(new_child)
|
|
new_child.parent.__children[child_idx] = None
|
|
new_child.__parent = self
|
|
self.__post_assign_children(new_children)
|
|
except Exception as exc_info:
|
|
# Reassign new children to their original parent
|
|
for child, idx_parent in current_new_children.items():
|
|
child_idx, parent = idx_parent
|
|
child.__parent = parent
|
|
parent.__children[child_idx] = child
|
|
for child in current_new_orphan:
|
|
child.__parent = None
|
|
|
|
# Reassign old children to self
|
|
self.__children = current_children
|
|
for child in current_children:
|
|
if child:
|
|
child.__parent = self
|
|
raise TreeError(exc_info)
|
|
|
|
@children.deleter
|
|
def children(self) -> None:
|
|
"""Delete child node(s)"""
|
|
for child in self.children:
|
|
if child is not None:
|
|
child.parent.__children.remove(child) # type: ignore
|
|
child.__parent = None
|
|
|
|
def __pre_assign_children(self: T, new_children: List[Optional[T]]) -> None:
|
|
"""Custom method to check before attaching children
|
|
Can be overridden with `_BinaryNode__pre_assign_children()`
|
|
|
|
Args:
|
|
new_children (List[Optional[Self]]): new children to be added
|
|
"""
|
|
pass
|
|
|
|
def __post_assign_children(self: T, new_children: List[Optional[T]]) -> None:
|
|
"""Custom method to check after attaching children
|
|
Can be overridden with `_BinaryNode__post_assign_children()`
|
|
|
|
Args:
|
|
new_children (List[Optional[Self]]): new children to be added
|
|
"""
|
|
pass
|
|
|
|
@property
|
|
def is_leaf(self) -> bool:
|
|
"""Get indicator if self is leaf node
|
|
|
|
Returns:
|
|
(bool)
|
|
"""
|
|
return not len([child for child in self.children if child])
|
|
|
|
def sort(self, **kwargs: Any) -> None:
|
|
"""Sort children, possible keyword arguments include ``key=lambda node: node.val``, ``reverse=True``
|
|
|
|
Examples:
|
|
>>> from bigtree import BinaryNode, print_tree
|
|
>>> a = BinaryNode(1)
|
|
>>> c = BinaryNode(3, parent=a)
|
|
>>> b = BinaryNode(2, parent=a)
|
|
>>> print_tree(a)
|
|
1
|
|
├── 3
|
|
└── 2
|
|
>>> a.sort(key=lambda node: node.val)
|
|
>>> print_tree(a)
|
|
1
|
|
├── 2
|
|
└── 3
|
|
"""
|
|
children = [child for child in self.children if child]
|
|
if len(children) == 2:
|
|
children.sort(**kwargs)
|
|
self.__children = children # type: ignore
|
|
|
|
def __repr__(self) -> str:
|
|
"""Print format of BinaryNode
|
|
|
|
Returns:
|
|
(str)
|
|
"""
|
|
class_name = self.__class__.__name__
|
|
node_dict = self.describe(exclude_prefix="_", exclude_attributes=[])
|
|
node_description = ", ".join([f"{k}={v}" for k, v in node_dict])
|
|
return f"{class_name}({node_description})"
|
|
|
|
|
|
T = TypeVar("T", bound=BinaryNode)
|