fork(1) download
  1. from typing import Optional
  2.  
  3. class TreeNode:
  4. def __init__(self, val=0, left=None, right=None):
  5. self.val = val
  6. self.left = left
  7. self.right = right
  8.  
  9. class Solution:
  10. def countNodes(self, root: Optional[TreeNode]) -> int:
  11. if not root: return 0
  12. def maxDepth(root):
  13. # this finds the maximum depth of the tree
  14. # it goes down just the left side of the tree, where nodes are guaranteed
  15. # O(d) - depth of left-side
  16. if not root:
  17. return -1
  18.  
  19. return 1 + maxDepth(root.left)
  20.  
  21.  
  22. d = maxDepth(root)
  23. def dfs(root, depth):
  24. if not root:
  25. return 0
  26. if depth == d:
  27. # if we've reached the bottom nodes, return 1 (counting itself)
  28. return 1
  29.  
  30. left, right, mid = 0, 0, (2 ** (d - depth)) // 2
  31. left = dfs(root.left, depth + 1)
  32. if left < mid:
  33. # determine if we must go to the right (halve the problem)
  34. # depends on how many bottom nodes are returned by the left recursion
  35. # for instance, if the depth of the tree is 4 (start from 0), and
  36. # the current depth is 2, then this subtree should have 4 bottom nodes
  37. # if it doesn't, the parent does not have to recurse rightward
  38. right = dfs(root.right, depth + 1)
  39.  
  40. return mid + right
  41.  
  42. bottom = dfs(root, 0)
  43. top = (2 ** (d)) - 1
  44. return top + bottom
  45.  
  46. def make_node(): return TreeNode(None, None, None)
  47. tree = make_node()
  48. tree.left = make_node()
  49. tree.right = make_node()
  50. tree.left.left = make_node()
  51.  
  52. print(Solution().countNodes(tree))
Success #stdin #stdout 0.04s 9652KB
stdin
Standard input is empty
stdout
6