diff --git a/tries/trie_find_height.py b/tries/trie_find_height.py index c37410f..618ff8a 100644 --- a/tries/trie_find_height.py +++ b/tries/trie_find_height.py @@ -15,25 +15,27 @@ def max_depth_recursive(root): if root is None: return 0 - if root.children == []: + if root.children: is None: return 1 - height = [max_depth_recursive(c) for c in root.children] - return max(height) + 1 + height = [max_depth_recursive(children) for children in root.children] + + return max(height) + 1 def max_depth_iterative(root): - stack = [] + stack, depth = [], 0 + if root is not None: stack.append((1, root)) - depth = 0 - - while stack != []: + while stack: this_depth, node = stack.pop() + if node is not None: + depth = max(depth, this_depth) for c in node.children: stack.append((this_depth + 1, c))