Count the number of nodes at a given level in a tree using DFS

Given an integer l and a tree represented as an undirected graph rooted at vertex 0. The task is to print the number of nodes present at level l.

Examples:

Input: l = 2 

Output: 4 

We have already discussed the BFS approach, in this post we will solve it using DFS.

Approach: The idea is to traverse the graph in a DFS manner. Take two variables, count and curr_level. Whenever the curr_level = l increment the value of the count.

Below is the implementation of the above approach: 

C++

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;

// Class to represent a graph
class Graph {

    // No. of vertices
    int V;

    // Pointer to an array containing
    // adjacency lists
    list<int>* adj;

    // A function used by NumOfNodes
    void DFS(vector<bool>& visited, int src, int& curr_level,
             int level, int& NumberOfNodes);

public:
    // Constructor
    Graph(int V);

    // Function to add an edge to graph
    void addEdge(int src, int des);

    // Returns the no. of nodes
    int NumOfNodes(int level);
};

Graph::Graph(int V)
{
    this->V = V;
    adj = new list<int>[V];
}

void Graph::addEdge(int src, int des)
{
    adj[src].push_back(des);
    adj[des].push_back(src);
}

// DFS function to keep track of
// number of nodes
void Graph::DFS(vector<bool>& visited, int src, int& curr_level,
                int level, int& NumberOfNodes)
{
    // Mark the current vertex as visited
    visited[src] = true;

    // If current level is equal
    // to the given level, increment
    // the no. of nodes
    if (level == curr_level) {
        NumberOfNodes++;
    }
    else if (level < curr_level)
        return;
    else {
        list<int>::iterator i;

        // Recur for the vertices
        // adjacent to the current vertex
        for (i = adj[src].begin(); i != adj[src].end(); i++) {
            if (!visited[*i]) {
                curr_level++;
                DFS(visited, *i, curr_level, level, NumberOfNodes);
            }
        }
    }
    curr_level--;
}

// Function to return the number of nodes
int Graph::NumOfNodes(int level)
{
    // To keep track of current level
    int curr_level = 0;

    // For keeping track of visited
    // nodes in DFS
    vector<bool> visited(V, false);

    // To store count of nodes at a
    // given level
    int NumberOfNodes = 0;

    DFS(visited, 0, curr_level, level, NumberOfNodes);

    return NumberOfNodes;
}

// Driver code
int main()
{
    int V = 8;

    Graph g(8);
    g.addEdge(0, 1);
    g.addEdge(0, 4);
    g.addEdge(0, 7);
    g.addEdge(4, 6);
    g.addEdge(4, 5);
    g.addEdge(4, 2);
    g.addEdge(7, 3);

    int level = 2;

    cout << g.NumOfNodes(level);

    return 0;
}

Java

import java.util.*;

// Class to represent a graph
class Graph {

    // No. of vertices
    private int V;

    // Pointer to an array containing
    // adjacency lists
    private List<Integer>[] adj;

    // A function used by NumOfNodes
    private void DFS(boolean[] visited, int src, int[] currLevel, int level, int[] numOfNodes) {
        // Mark the current vertex as visited
        visited[src] = true;

        // If current level is equal
        // to the given level, increment
        // the no. of nodes
        if (level == currLevel[0]) {
            numOfNodes[0]++;
        } else if (level < currLevel[0]) {
            return;
        } else {
            // Recur for the vertices
            // adjacent to the current vertex
            for (int neighbor : adj[src]) {
                if (!visited[neighbor]) {
                    currLevel[0]++;
                    DFS(visited, neighbor, currLevel, level, numOfNodes);
                }
            }
        }
        currLevel[0]--;
    }

    // Constructor
    public Graph(int V) {
        this.V = V;
        adj = new ArrayList[V];
        for (int i = 0; i < V; i++) {
            adj[i] = new ArrayList<>();
        }
    }

    // Function to add an edge to the graph
    public void addEdge(int src, int des) {
        adj[src].add(des);
        adj[des].add(src);
    }

    // Function to return the number of nodes at a given level
    public int numOfNodes(int level) {
        // To keep track of the current level
        int[] currLevel = { 0 };

        // For keeping track of visited
        // nodes in DFS
        boolean[] visited = new boolean[V];

        // To store count of nodes at a
        // given level
        int[] numOfNodes = { 0 };

        DFS(visited, 0, currLevel, level, numOfNodes);

        return numOfNodes[0];
    }
}

public class Main {

    // Driver code
    public static void main(String[] args) {
        int V = 8;

        Graph g = new Graph(V);
        g.addEdge(0, 1);
        g.addEdge(0, 4);
        g.addEdge(0, 7);
        g.addEdge(4, 6);
        g.addEdge(4, 5);
        g.addEdge(4, 2);
        g.addEdge(7, 3);

        int level = 2;

        System.out.println(g.numOfNodes(level));
    }
}

Python3

# Python3 implementation of the approach

# Class to represent a graph
class Graph:

    def __init__(self, V):

        # No. of vertices
        self.V = V

        # Pointer to an array containing
        # adjacency lists
        self.adj = [[] for i in range(self.V)]

    def addEdge(self, src, des):

        self.adj[src].append(des)
        self.adj[des].append(src)

    # DFS function to keep track of
    # number of nodes
    def DFS(self, visited, src, curr_level,
            level, NumberOfNodes):

        # Mark the current vertex as visited
        visited[src] = True

        # If current level is equal
        # to the given level, increment
        # the no. of nodes
        if (level == curr_level):
            NumberOfNodes += 1

        elif (level < curr_level):
            return
        else:

            # Recur for the vertices
            # adjacent to the current vertex
            for i in self.adj[src]:

                if (not visited[i]):
                    curr_level += 1
                    curr_level, NumberOfNodes = self.DFS(
                        visited, i, curr_level,
                        level, NumberOfNodes)

        curr_level -= 1

        return curr_level, NumberOfNodes

    # Function to return the number of nodes
    def NumOfNodes(self, level):

        # To keep track of current level
        curr_level = 0

        # For keeping track of visited
        # nodes in DFS
        visited = [False for i in range(self.V)]

        # To store count of nodes at a
        # given level
        NumberOfNodes = 0

        curr_level, NumberOfNodes = self.DFS(
            visited, 0, curr_level,
            level, NumberOfNodes)

        return NumberOfNodes

# Driver code
if __name__=='__main__':

    V = 8

    g = Graph(8)
    g.addEdge(0, 1)
    g.addEdge(0, 4)
    g.addEdge(0, 7)
    g.addEdge(4, 6)
    g.addEdge(4, 5)
    g.addEdge(4, 2)
    g.addEdge(7, 3)

    level = 2

    print(g.NumOfNodes(level))

# This code is contributed by pratham76

C#

using System;
using System.Collections.Generic;

// Class to represent a graph
class Graph
{
    // No. of vertices
    int V;

    // Pointer to an array containing adjacency lists
    List<int>[] adj;

    // A function used by NumOfNodes
    void DFS(List<bool> visited, int src, ref int curr_level,
             int level, ref int NumberOfNodes)
    {
        // Mark the current vertex as visited
        visited[src] = true;

        // If current level is equal to the given level, increment the no. of nodes
        if (level == curr_level)
        {
            NumberOfNodes++;
        }
        else if (level < curr_level)
        {
            return;
        }
        else
        {
            List<int>.Enumerator i;

            // Recur for the vertices adjacent to the current vertex
            i = adj[src].GetEnumerator();
            while (i.MoveNext())
            {
                if (!visited[i.Current])
                {
                    curr_level++;
                    DFS(visited, i.Current, ref curr_level, level, ref NumberOfNodes);
                }
            }
        }
        curr_level--;
    }

    public Graph(int V)
    {
        this.V = V;
        adj = new List<int>[V];
        for (int i = 0; i < V; ++i)
        {
            adj[i] = new List<int>();
        }
    }

    public void addEdge(int src, int des)
    {
        adj[src].Add(des);
        adj[des].Add(src);
    }

    public int NumOfNodes(int level)
    {
        // To keep track of current level
        int curr_level = 0;

        // For keeping track of visited nodes in DFS
        List<bool> visited = new List<bool>(V);
        for (int i = 0; i < V; ++i)
        {
            visited.Add(false);
        }

        // To store count of nodes at a given level
        int NumberOfNodes = 0;

        DFS(visited, 0, ref curr_level, level, ref NumberOfNodes);

        return NumberOfNodes;
    }
}

class MainClass
{
    public static void Main()
    {
        int V = 8;

        Graph g = new Graph(8);
        g.addEdge(0, 1);
        g.addEdge(0, 4);
        g.addEdge(0, 7);
        g.addEdge(4, 6);
        g.addEdge(4, 5);
        g.addEdge(4, 2);
        g.addEdge(7, 3);

        int level = 2;

        Console.WriteLine(g.NumOfNodes(level));
    }
}

// This code is contributed by Prince Kumar

Javascript

// JavaScript implementation of the approach

// Class to represent a graph
class Graph {
constructor(V) {
// No. of vertices
this.V = V;

// Pointer to an array containing adjacency lists
this.adj = Array.from({ length: this.V }, () => []);
}

addEdge(src, des) {
this.adj[src].push(des);
this.adj[des].push(src);
}

// DFS function to keep track of number of nodes
DFS(visited, src, curr_level, level, NumberOfNodes) {

// Mark the current vertex as visited
visited[src] = true;


// If current level is equal to the given level, increment the no. of nodes
if (level == curr_level) {
  NumberOfNodes += 1;
} else if (level < curr_level) {
  return;
} else {

  // Recur for the vertices adjacent to the current vertex
  for (const i of this.adj[src]) {
    if (!visited[i]) {
      curr_level += 1;
      [curr_level, NumberOfNodes] = this.DFS(
        visited,
        i,
        curr_level,
        level,
        NumberOfNodes
      );
    }
  }
}
curr_level -= 1;
return [curr_level, NumberOfNodes];
}

// Function to return the number of nodes
NumOfNodes(level)
{

// To keep track of current level
let curr_level = 0;


// For keeping track of visited nodes in DFS
let visited = new Array(this.V).fill(false);

// To store count of nodes at a given level
let NumberOfNodes = 0;

[curr_level, NumberOfNodes] = this.DFS(
  visited,
  0,
  curr_level,
  level,
  NumberOfNodes
);

return NumberOfNodes;
}
}

// Driver code
const g = new Graph(8);
g.addEdge(0, 1);
g.addEdge(0, 4);
g.addEdge(0, 7);
g.addEdge(4, 6);
g.addEdge(4, 5);
g.addEdge(4, 2);
g.addEdge(7, 3);

const level = 2;
console.log(g.NumOfNodes(level));

// This code is contributed by lokeshpotta20.

Output

4

Complexity Analysis:

• Time Complexity : O(N), where N is the total number of nodes in the graph.
• Auxiliary Space: O(N)