Introduction

To count how many nodes lie within distance k of a given node in a rooted tree, treat the tree as an undirected graph and run a breadth-first search from the target. The fact that the tree is rooted matters for storage and traversal order, but graph distance still counts edges in both directions. That means a correct solution must be able to move from a node to its parent as well as to its children.

Distance in a Rooted Tree Is Still Graph Distance

A rooted tree is often stored with parent and child relationships, which makes it easy to think only in the downward direction. That is the main source of bugs in this problem.

Suppose the target node is somewhere in the middle of the tree. Nodes within distance two may include:

• its parent
• its grandparent
• its siblings through the parent
• its children and grandchildren

A traversal that only follows children misses half of the valid search space. So the first step is to create a representation that allows upward and downward movement.

Build an Undirected Adjacency List

If the tree is stored as parent pointers, convert it to an adjacency list once. That gives a clean graph view for BFS.

1from collections import defaultdict
2
3
4def build_adjacency(parent):
5    adj = defaultdict(list)
6    for child in range(1, len(parent)):
7        p = parent[child]
8        adj[p].append(child)
9        adj[child].append(p)
10    return adj
11
12
13parent = [-1, 0, 0, 1, 1, 3]
14adj = build_adjacency(parent)
15print(dict(adj))

If your input already comes as edges, the same idea applies. Add both directions for every edge because distance queries need an undirected view.

Use BFS to Count Nodes Up to Distance k

Breadth-first search is the natural algorithm here because it explores nodes in increasing distance order. Once a node is reached at distance k, you count it but stop expanding from it.

1from collections import deque
2
3
4def count_within_distance(adj, start, k):
5    visited = set([start])
6    queue = deque([(start, 0)])
7    count = 0
8
9    while queue:
10        node, dist = queue.popleft()
11        count += 1
12
13        if dist == k:
14            continue
15
16        for neighbor in adj[node]:
17            if neighbor not in visited:
18                visited.add(neighbor)
19                queue.append((neighbor, dist + 1))
20
21    return count
22
23
24parent = [-1, 0, 0, 1, 1, 3]
25adj = build_adjacency(parent)
26print(count_within_distance(adj, start=1, k=2))

This counts the starting node itself because its distance is zero. In most mathematical formulations, “within distance k” includes the source node.

Count Exactly Distance k If Needed

Some problem statements ask for nodes at exactly distance k, not at most k. In that case, the BFS is almost the same, but the counting rule changes.

1def count_exactly_distance(adj, start, k):
2    visited = set([start])
3    queue = deque([(start, 0)])
4    count = 0
5
6    while queue:
7        node, dist = queue.popleft()
8
9        if dist == k:
10            count += 1
11            continue
12
13        for neighbor in adj[node]:
14            if neighbor not in visited:
15                visited.add(neighbor)
16                queue.append((neighbor, dist + 1))
17
18    return count

This difference looks small, but it changes the result enough that you should decide it explicitly before implementing.

Complexity and When BFS Is Enough

For one query, BFS visits each node at most once, so the time complexity is O(n) in the worst case. Space complexity is also O(n) because of the queue and visited set.

That is the right baseline for most tasks. Only move to advanced preprocessing, such as subtree statistics, centroid decomposition, or LCA-based structures, if you have many queries on a large static tree and measurements show BFS is the bottleneck.

In other words, do not overengineer the first solution. A clear O(n) BFS is often exactly what the problem needs.

Subtree Information Alone Is Not Sufficient

People sometimes try to solve this using only subtree sizes because the input is a rooted tree. That works only for special cases such as counting descendants within a depth limit from the root or from an ancestor. It does not solve the general “within distance of any target node” problem, because valid nodes may lie outside the target’s subtree.

This is why the undirected representation matters so much.

Common Pitfalls

The most common bug is traversing only downward from the target. That misses ancestors and nodes reached through ancestors.

Another mistake is forgetting the visited set after turning the tree into an undirected graph. A tree has no cycles in directed parent-child form, but the undirected version does let you walk back and forth unless you mark visited nodes.

Developers also mix up “within distance k” and “exactly distance k.” Decide that rule up front and write the counting logic accordingly.

Summary

• For distance queries, treat the rooted tree as an undirected graph.
• Build parent links or an adjacency list so traversal can move both up and down.
• Use BFS from the target node and stop expanding when distance k is reached.
• Decide explicitly whether the problem means distance at most k or exactly k.
• A simple O(n) BFS is usually the correct starting solution for one or a few queries.