From 778e2010d6ae89c61a93672e49b86041b6ca1108 Mon Sep 17 00:00:00 2001 From: Vinayak Upadhyay Date: Mon, 16 Oct 2023 22:16:44 +0530 Subject: [PATCH] Added functionality to calculate the diameter of given binary tree (#10526) * Added code to find diameter of given binary tree * Modified diameter_of_binary_tree file * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update diameter_of_binary_tree.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update diameter_of_binary_tree.py --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Christian Clauss --- .../binary_tree/diameter_of_binary_tree.py | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 data_structures/binary_tree/diameter_of_binary_tree.py diff --git a/data_structures/binary_tree/diameter_of_binary_tree.py b/data_structures/binary_tree/diameter_of_binary_tree.py new file mode 100644 index 000000000..bbe70b028 --- /dev/null +++ b/data_structures/binary_tree/diameter_of_binary_tree.py @@ -0,0 +1,72 @@ +""" +The diameter/width of a tree is defined as the number of nodes on the longest path +between two end nodes. +""" +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class Node: + data: int + left: Node | None = None + right: Node | None = None + + def depth(self) -> int: + """ + >>> root = Node(1) + >>> root.depth() + 1 + >>> root.left = Node(2) + >>> root.depth() + 2 + >>> root.left.depth() + 1 + >>> root.right = Node(3) + >>> root.depth() + 2 + """ + left_depth = self.left.depth() if self.left else 0 + right_depth = self.right.depth() if self.right else 0 + return max(left_depth, right_depth) + 1 + + def diameter(self) -> int: + """ + >>> root = Node(1) + >>> root.diameter() + 1 + >>> root.left = Node(2) + >>> root.diameter() + 2 + >>> root.left.diameter() + 1 + >>> root.right = Node(3) + >>> root.diameter() + 3 + """ + left_depth = self.left.depth() if self.left else 0 + right_depth = self.right.depth() if self.right else 0 + return left_depth + right_depth + 1 + + +if __name__ == "__main__": + from doctest import testmod + + testmod() + root = Node(1) + root.left = Node(2) + root.right = Node(3) + root.left.left = Node(4) + root.left.right = Node(5) + r""" + Constructed binary tree is + 1 + / \ + 2 3 + / \ + 4 5 + """ + print(f"{root.diameter() = }") # 4 + print(f"{root.left.diameter() = }") # 3 + print(f"{root.right.diameter() = }") # 1