Count of nodes at a distance K for every node

Given a tree of N vertices and a positive integer K. The task is to find the count of nodes at a distance K for every node.

Examples:

Input: N = 5, K = 2, edges[][] = {{1, 2}, {1, 3}, {2, 4}, {2, 5}}
Output: 2 1 1 2 2
Explanation:

• From node 1, there are 2 nodes that are at a distance of 2: node 4 and node 5.
• From node 2, there is 1 node that is at a distance of 2: node 3.
• From node 3, there is 1 node that is at a distance of 2: node 2.
• From node 4, there are 2 nodes that are at a distance of 2: node 1 and node 5.
• From node 5, there are 2 nodes that are at a distance of 2: node 1 and node 4.

Input: N = 5 K = 3, edges[][] = {{1, 2}, {2, 3}, {3, 4}, {4, 5}}
Output: 1 1 0 1 1

Approach: To solve the problem, follow the below idea:

The approach involves a depth-first search (DFS) traversal of the tree starting from each node to efficiently count the number of nodes at a distance K from all nodes.

Step-by-step algorithm:

• Perform a DFS traversal from a given node v. It keeps track of the current distance d from the starting node and the target distance K to find nodes at distance k from the starting node.
• During the traversal, it marks each visited node and updates the distance of each node from the starting node.
• If the current distance d equals the target distance k, it increments the count of nodes at distance k (ans[v]++).
• If the current distance d exceeds the target distance k, it returns without further exploration, as nodes beyond distance k are not considered.
• For each adjacent node u of the current node v, if u has not been visited yet, the dfs function is called recursively to explore u with an incremented distance d+1.
• After performing the DFS traversal from each node, output the count of nodes at distance k for each node in the tree.

Below is the implementation of the above approach:

C++

#include <bits/stdc++.h>
using namespace std;

// Maximum number of vertices in the tree
const int MAX_N = 100005;
// Adjacency list representation of the tree
vector<int> adj[MAX_N];
// Array to mark visited nodes
bool visited[MAX_N];
// Array to store distances from the starting node
int dist[MAX_N];
// Array to store the count of nodes at distance K
int ans[MAX_N];

void dfs(int startNode, int currentNode, int distance,
         int K)
{
    visited[currentNode] = true;
    dist[currentNode] = distance;
    if (distance == K) {
        ans[startNode]++;
    }
    if (distance > K) {
        return;
    }
    for (int u : adj[currentNode]) {
        if (!visited[u]) {
            dfs(startNode, u, distance + 1, K);
        }
    }
}

int main()
{
    int N = 5;
    int K = 2;

    vector<vector<int>> edges{
        { 1, 2 }, { 1, 3 }, { 2, 4 }, { 2, 5 }
    };

    for (int i = 0; i < N - 1; i++) {
        adj[edges[i][0]].push_back(edges[i][1]);
        adj[edges[i][1]].push_back(edges[i][0]);
    }

    // Perform DFS traversal from each node
    for (int i = 1; i <= N; i++) {
        fill(visited, visited + N + 1, false);
        dfs(i, i, 0, K);
    }

    for (int i = 1; i <= N; i++) {
        cout << i << " " << ans[i] << endl;
    }
    return 0;
}

Java

import java.util.*;

public class Main {
    // Maximum number of vertices in the tree
    static final int MAX_N = 100005;
    // Adjacency list representation of the tree
    static List<Integer>[] adj = new ArrayList[MAX_N];
    // Array to mark visited nodes
    static boolean[] visited = new boolean[MAX_N];
    // Array to store distances from the starting node
    static int[] dist = new int[MAX_N];
    // Array to store the count of nodes at distance K
    static int[] ans = new int[MAX_N];

    public static void dfs(int startNode, int currentNode, int distance, int K) {
        visited[currentNode] = true;
        dist[currentNode] = distance;
        if (distance == K) {
            ans[startNode]++;
        }
        if (distance > K) {
            return;
        }
        for (int u : adj[currentNode]) {
            if (!visited[u]) {
                dfs(startNode, u, distance + 1, K);
            }
        }
    }

    public static void main(String[] args) {
        int N = 5;
        int K = 2;

        // Initializing adjacency list
        for (int i = 1; i <= N; i++) {
            adj[i] = new ArrayList<>();
        }

        // Edge list
        int[][] edges = {
            { 1, 2 }, { 1, 3 }, { 2, 4 }, { 2, 5 }
        };

        // Adding edges to the adjacency list
        for (int i = 0; i < N - 1; i++) {
            adj[edges[i][0]].add(edges[i][1]);
            adj[edges[i][1]].add(edges[i][0]);
        }

        // Perform DFS traversal from each node
        for (int i = 1; i <= N; i++) {
            Arrays.fill(visited, false);
            dfs(i, i, 0, K);
        }

        // Print the count of nodes at distance K for each node
        for (int i = 1; i <= N; i++) {
            System.out.println(i + " " + ans[i]);
        }
    }
}

Python3

from collections import defaultdict

# Maximum number of vertices in the tree
MAX_N = 100005
# Adjacency list representation of the tree
adj = defaultdict(list)
# Array to store the count of nodes at distance K
ans = [0] * MAX_N

def dfs(current_node, parent, distance, K):
    if distance == K:
        ans[current_node] += 1
        return
    for u in adj[current_node]:
        if u != parent:
            dfs(u, current_node, distance + 1, K)

def main():
    N = 5
    K = 2

    edges = [
        [1, 2], [1, 3], [2, 4], [2, 5]
    ]

    for edge in edges:
        u, v = edge
        adj[u].append(v)
        adj[v].append(u)

    # Perform DFS traversal from each node
    for i in range(1, N + 1):
        dfs(i, -1, 0, K)

    print("Node    Count of nodes at distance", K)
    for i in range(1, N + 1):
        print(i, ans[i])

if __name__ == "__main__":
    main()


# This code is contributed by shivamgupta0987654321

JavaScript

// Maximum number of vertices in the tree
const MAX_N = 100005;
// Adjacency list representation of the tree
const adj = new Array(MAX_N).fill(0).map(() => []);
// Array to mark visited nodes
const visited = new Array(MAX_N).fill(false);
// Array to store distances from the starting node
const dist = new Array(MAX_N).fill(0);
// Array to store the count of nodes at distance K
const ans = new Array(MAX_N).fill(0);

const dfs = (startNode, currentNode, distance, K) => {
  visited[currentNode] = true;
  dist[currentNode] = distance;
  if (distance === K) {
    ans[startNode]++;
  }
  if (distance > K) {
    return;
  }
  for (let u of adj[currentNode]) {
    if (!visited[u]) {
      dfs(startNode, u, distance + 1, K);
    }
  }
};

const main = () => {
  const N = 5;
  const K = 2;

  const edges = [
    [1, 2],
    [1, 3],
    [2, 4],
    [2, 5],
  ];

  for (let i = 0; i < N - 1; i++) {
    adj[edges[i][0]].push(edges[i][1]);
    adj[edges[i][1]].push(edges[i][0]);
  }

  // Perform DFS traversal from each node
  for (let i = 1; i <= N; i++) {
    visited.fill(false);
    dfs(i, i, 0, K);
  }

  for (let i = 1; i <= N; i++) {
    console.log(`${i} ${ans[i]}`);
  }
};

main();

Output

1 2
2 1
3 1
4 2
5 2

Time complexity: O(N * K), where N is the number of nodes and K is the distance between the nodes.
Auxiliary Space: O(N * K)