feat(datastructures, trees): binary tree branch sum and node depth

Author: BrianLusina

datastructures/trees/binary/node.py

@@ -17,7 +17,7 @@ def __init__(
         right: Optional["BinaryTreeNode"] = None,
         key: Optional[Any] = None,
         parent: Optional["BinaryTreeNode"] = None,
-        next: Optional["BinaryTreeNode"] = None,
+        nxt: Optional["BinaryTreeNode"] = None,
     ) -> None:
         """
         Constructor for BinaryTreeNode class. This will create a new node with the provided data and optional
@@ -30,7 +30,7 @@ def __init__(
             right (Optional[BinaryTreeNode]): Right child of the node
             key (Optional[Any]): Key for the node, if not provided a hash of the data is used
             parent (Optional[BinaryTreeNode]): Parent of the node
-            next (Optional[BinaryTreeNode]): Next child of the node which is the sibling of the node. The sibling is the
+            nxt (Optional[BinaryTreeNode]): Next child of the node which is the sibling of the node. The sibling is the
             node on the same level as this node. If this is the rightmost node in the tree that is not on the last level
             of the tree, then this is the next node on the next level starting from the left. If this is the last node
             in the tree, then this is None.
@@ -42,7 +42,7 @@ def __init__(
         # node on a given level, then it is connected to the first node on the next level. If this node is the last node
         # in the tree on the last node, then it is pointed to None. By default, it is set to None.
         # Note that if this is the root node, it is connected to the left most node on the next level.
-        self.next: Optional[BinaryTreeNode] = next
+        self.next: Optional[BinaryTreeNode] = nxt
 
     def insert_node(self, data: T) -> None:
         """

datastructures/trees/binary/tree/binary_tree.py

@@ -960,3 +960,144 @@ def longest_uni_value_path(self) -> int:
         The length of the path between two nodes is represented by the number of edges between them.
         """
         return longest_uni_value_path(self.root)
+
+    def branch_sum(self) -> List[Any]:
+        """
+        Retrieves the sum of branches ordered from left most branch sum to right most branch sum. A branch sum is the
+        sum of all values in a binary tree branch. A binary tree branch is a path of nodes in a tree that starts at the
+        root node and ends at any leaf node.
+
+        For example given the tree below
+
+                   1
+                /     \
+               2       3
+             /   \    /  \
+            4     5  6    7
+          /   \  /
+         8    9 10
+
+        The output should be:
+        [15, 16, 18, 10, 11]
+
+        15 == 1 + 2 + 4 + 8
+        16 == 1 + 2 + 4 + 9
+        18 == 1 + 2 + 5 + 10
+        10 == 1 + 3 + 6
+        11 == 1 + 3 + 7
+
+        The time complexity for the solution is O(n) because we have to traverse each element in the tree and the space
+        complexity is also O(n) because of the recursion.
+        """
+        if not self.root:
+            return []
+
+        branch_totals = []
+
+        def preorder_traversal(
+            node: Optional[BinaryTreeNode], running_sum: Any
+        ) -> None:
+            nonlocal branch_totals
+            if not node:
+                return
+
+            total_sum = running_sum + node.data
+            if not node.left and not node.right:
+                branch_totals.append(total_sum)
+
+            preorder_traversal(node.left, total_sum)
+            preorder_traversal(node.right, total_sum)
+
+        preorder_traversal(self.root, 0)
+        return branch_totals
+
+    def sum_of_node_depths(self) -> int:
+        """
+        Returns the sum of depths of all nodes in the binary tree.
+        Node Depth is the distance between a node in a binary tree and the tree's root.
+        For example given the tree below
+
+                   1
+                /     \
+               2       3
+             /   \    /  \
+            4     5  6    7
+          /   \
+         8    9
+        Output should be: 16
+        The depth of node with value 2 is 1
+        The depth of node with value 3 is 1
+        The depth of node with value 4 is 2
+        The depth of node with value 5 is 2
+        etc ...
+        Summing all of these depths yields 16
+
+        This uses an iterative approach with a stack to get the sum of all the node depths. The stack will store the
+        node alongside the depth of the node and be used to calculate the total depth. This incurs a time complexity
+        cost of O(n) as we traverse all the nodes of the tree and space complexity cost of O(h) where h is the height
+        of the tree as we store all the nodes in the stack.
+
+        Returns:
+            int: Sum of depths of all nodes in the binary tree.
+        """
+
+        if not self.root:
+            return 0
+
+        stack: List[Tuple[BinaryTreeNode, int]] = [(self.root, 0)]
+        total_depth: int = 0
+
+        while stack:
+            node, depth = stack.pop()
+            total_depth += depth
+            if node.left:
+                stack.append((node.left, depth + 1))
+            if node.right:
+                stack.append((node.right, depth + 1))
+
+        return total_depth
+
+    def sum_of_node_depths_recursive(self) -> int:
+        """
+        Returns the sum of depths of all nodes in the binary tree using recursion.
+        Node Depth is the distance between a node in a binary tree and the tree's root.
+        For example given the tree below
+
+                   1
+                /     \
+               2       3
+             /   \    /  \
+            4     5  6    7
+          /   \
+         8    9
+        Output should be: 16
+        The depth of node with value 2 is 1
+        The depth of node with value 3 is 1
+        The depth of node with value 4 is 2
+        The depth of node with value 5 is 2
+        etc ...
+        Summing all of these depths yields 16
+
+        This uses a recursive approach to get the sum of all the node depths. This incurs a time complexity
+        cost of O(n) as we traverse all the nodes of the tree and space complexity cost of O(h) where h is the height
+        of the tree as we call the recursion stack h times on the height of the tree.
+
+        Returns:
+            int: Sum of depths of all nodes in the binary tree.
+        """
+
+        if not self.root:
+            return 0
+
+        def sum_of_node_depths_helper(
+            node: Optional[BinaryTreeNode], depth: int
+        ) -> int:
+            if not node:
+                return 0
+            return (
+                depth
+                + sum_of_node_depths_helper(node.left, depth)
+                + sum_of_node_depths_helper(node.right, depth)
+            )
+
+        return sum_of_node_depths_helper(self.root, 0)

datastructures/trees/heaps/__init__.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, List
+from typing import Any, List, Optional
 
 
 class Heap(ABC):
@@ -8,12 +8,19 @@ class Heap(ABC):
     """
 
     @abstractmethod
-    def insert_data(self, data: Any):
+    def insert(self, data: Any):
         """
         Inserts a data value into the heap
         """
         raise NotImplementedError("Not yet implemented")
 
+    @abstractmethod
+    def peak(self) -> Optional[Any]:
+        """
+        Returns data at the top of the heap without removing it
+        """
+        raise NotImplementedError("Not yet implemented")
+
     @abstractmethod
     def delete(self) -> Any:
         """
@@ -38,37 +45,46 @@ def get_right_child_index(i: int) -> int:
 
 class ArrayBasedHeap(Heap):
     """
-    Heap datastructure that uses an array as the underlying datastructure to build a heap.
+    Heap data structure that uses an array as the underlying data structure to build a heap.
     """
 
-    def __init__(self):
+    def __init__(self, data: Optional[List[Any]] = None):
         super().__init__()
-        self.data: List[Any] = []
+        if data is None:
+            data = []
+        self.heap: List[Any] = data
 
     def __len__(self):
-        return len(self.data)
+        return len(self.heap)
 
     @property
-    def root_node(self):
+    def root_node(self) -> Optional[Any]:
+        """
+        Retrieves the root node of the Heap
+        """
+        if len(self.heap) == 0:
+            raise Exception("Heap is empty")
+        return self.heap[0]
+
+    def peak(self) -> Optional[Any]:
         """
         Retrieves the root node of the Heap
-        :return:
         """
-        if len(self.data) == 0:
+        if len(self.heap) == 0:
             raise Exception("Heap is empty")
-        return self.data[0]
+        return self.heap[0]
 
     @property
     def last_node(self):
         """
         Returns the last node of the heap
         :return:
         """
-        if len(self.data) == 0:
+        if len(self.heap) == 0:
             raise Exception("Heap is empty")
-        return self.data[len(self.data) - 1]
+        return self.heap[len(self.heap) - 1]
 
-    def insert_data(self, data: Any):
+    def insert(self, data: Any):
         """
         Inserts a value into the heap
         :param data: element to insert into the heap

datastructures/trees/heaps/binary/max_heap/__init__.py

@@ -91,7 +91,7 @@ def __build_heap(self, array: List[HeapNode]) -> List[HeapNode]:
 
         return array
 
-    def insert_data(self, data: Any):
+    def insert(self, data: Any):
         """
         Inserts a value into the heap
         """

datastructures/trees/heaps/binary/max_heap/max_array_heap.py

@@ -11,42 +11,42 @@ class MaxArrayBasedHeap(ArrayBasedHeap):
     def __init__(self):
         super().__init__()
 
-    def insert_data(self, value: Any):
+    def insert(self, value: Any):
         """
         Inserts a value into the heap
         """
         # add the value into the last node
-        self.data.append(value)
+        self.heap.append(value)
 
         # keep track of the index of the newly inserted node
-        new_node_index = len(self.data) - 1
+        new_node_index = len(self.heap) - 1
 
         # the following executes the "trickle up" algorithm. If the new node is not in the root position and it's greater
         # than its parent node
         while (
             new_node_index > 0
-            and self.data[new_node_index]
-            > self.data[self.get_parent_index(new_node_index)]
+            and self.heap[new_node_index]
+            > self.heap[self.get_parent_index(new_node_index)]
         ):
             # swap the new node with the parent node
             (
-                self.data[self.get_parent_index(new_node_index)],
-                self.data[new_node_index],
+                self.heap[self.get_parent_index(new_node_index)],
+                self.heap[new_node_index],
             ) = (
-                self.data[new_node_index],
-                self.data[self.get_parent_index(new_node_index)],
+                self.heap[new_node_index],
+                self.heap[self.get_parent_index(new_node_index)],
             )
 
             # update the index of the new node
             new_node_index = self.get_parent_index(new_node_index)
 
     def delete(self) -> Any:
-        if len(self.data) == 0:
+        if len(self.heap) == 0:
             raise Exception("Heap is empty")
 
         # we only ever delete the root node from a heap, so we pop the last node from the array and make it the root node
-        root_node = self.data[0]
-        self.data[0] = self.data.pop()
+        root_node = self.heap[0]
+        self.heap[0] = self.heap.pop()
 
         # track the current index of the "trickle node". This is the node that will be moved into the correct position
         trickle_node_index = 0
@@ -57,9 +57,9 @@ def delete(self) -> Any:
             larger_child_index = self.__calculate_larger_child_index(trickle_node_index)
 
             # swap the trickle node with its larger child
-            self.data[trickle_node_index], self.data[larger_child_index] = (
-                self.data[larger_child_index],
-                self.data[trickle_node_index],
+            self.heap[trickle_node_index], self.heap[larger_child_index] = (
+                self.heap[larger_child_index],
+                self.heap[trickle_node_index],
             )
 
             trickle_node_index = larger_child_index
@@ -76,20 +76,20 @@ def __has_greater_child(self, index: int) -> bool:
         left_child_index = self.get_left_child_index(index)
         right_child_index = self.get_right_child_index(index)
 
-        left_child_exists = left_child_index < len(self.data)
-        right_child_exists = right_child_index < len(self.data)
+        left_child_exists = left_child_index < len(self.heap)
+        right_child_exists = right_child_index < len(self.heap)
 
         if left_child_exists and right_child_exists:
-            left_child = self.data[left_child_index]
-            right_child = self.data[right_child_index]
+            left_child = self.heap[left_child_index]
+            right_child = self.heap[right_child_index]
 
-            return left_child > self.data[index] or right_child > self.data[index]
+            return left_child > self.heap[index] or right_child > self.heap[index]
         elif left_child_exists and not right_child_exists:
-            left_child = self.data[left_child_index]
-            return left_child > self.data[index]
+            left_child = self.heap[left_child_index]
+            return left_child > self.heap[index]
         elif right_child_exists and not left_child_exists:
-            right_child = self.data[right_child_index]
-            return right_child > self.data[index]
+            right_child = self.heap[right_child_index]
+            return right_child > self.heap[index]
         else:
             return False
 
@@ -100,14 +100,14 @@ def __calculate_larger_child_index(self, index: int) -> int:
         :return: The position of the larger child
         """
         # if there is no right child
-        if not self.data[self.get_right_child_index(index)]:
+        if not self.heap[self.get_right_child_index(index)]:
             # return the left child index
             return self.get_left_child_index(index)
 
         # if right child value is greater than left child value
         if (
-            self.data[self.get_right_child_index(index)]
-            > self.data[self.get_left_child_index(index)]
+            self.heap[self.get_right_child_index(index)]
+            > self.heap[self.get_left_child_index(index)]
         ):
             # return the right child index
             return self.get_right_child_index(index)