Exercise 2: Count All Nodes 🔢

Your Task

Write a function count_nodes(root) that returns the total number of nodes in a binary tree.

---

Examples

#       1
#      / \
#     2   3
#    / \   \
#   4   5   6

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.right = TreeNode(6)

print(count_nodes(root))  # 6

#   1

print(count_nodes(TreeNode(1)))  # 1

print(count_nodes(None))  # 0

---

Expected Complexity

	Complexity
Time	O(n) — must visit every node
Space	O(h) — recursion stack depth

---

Approach

1. Base case: if root is None, return 0 (no nodes here)
2. Count nodes in the left subtree
3. Count nodes in the right subtree
4. Return 1 + left_count + right_count (the 1 is for the current node)

---

Hints