|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import copy |
| 4 | +import heapq |
4 | 5 | from typing import Any, Dict, Generator, Iterable, List, Optional, Set, Tuple, TypeVar |
5 | 6 |
|
6 | 7 | from bigtree.globals import ASSERTIONS |
@@ -92,8 +93,9 @@ class BaseNode: |
92 | 93 | 2. ``is_root``: Get indicator if self is root node |
93 | 94 | 3. ``is_leaf``: Get indicator if self is leaf node |
94 | 95 | 4. ``root``: Get root node of tree |
95 | | - 5. ``depth``: Get depth of self |
96 | | - 6. ``max_depth``: Get maximum depth from root to leaf node |
| 96 | + 5. ``diameter``: Get diameter of self |
| 97 | + 6. ``depth``: Get depth of self |
| 98 | + 7. ``max_depth``: Get maximum depth from root to leaf node |
97 | 99 |
|
98 | 100 | **BaseNode Methods** |
99 | 101 |
|
@@ -502,6 +504,37 @@ def root(self: T) -> T: |
502 | 504 | return self |
503 | 505 | return self.parent.root |
504 | 506 |
|
| 507 | + @property |
| 508 | + def diameter(self) -> int: |
| 509 | + """Get diameter of tree or subtree, the length of longest path between any two nodes |
| 510 | +
|
| 511 | + Returns: |
| 512 | + (int) |
| 513 | + """ |
| 514 | + diameter = 0 |
| 515 | + |
| 516 | + if self.is_leaf: |
| 517 | + return diameter |
| 518 | + |
| 519 | + def _recursive_diameter(node: T) -> int: |
| 520 | + """Recursively iterate through node and its children to get diameter of node |
| 521 | +
|
| 522 | + Args: |
| 523 | + node (Node): current node |
| 524 | +
|
| 525 | + Returns: |
| 526 | + (int) |
| 527 | + """ |
| 528 | + nonlocal diameter |
| 529 | + if node.is_leaf: |
| 530 | + return 1 |
| 531 | + child_length = [_recursive_diameter(child) for child in node.children] |
| 532 | + diameter = max(diameter, sum(heapq.nlargest(2, child_length))) |
| 533 | + return 1 + max(child_length) |
| 534 | + |
| 535 | + _recursive_diameter(self) |
| 536 | + return diameter |
| 537 | + |
505 | 538 | @property |
506 | 539 | def depth(self) -> int: |
507 | 540 | """Get depth of self, indexing starts from 1 |
|
0 commit comments