Count pairs of nodes having minimum distance between them equal to the difference of their distances from root

Given an N-ary Tree consisting of N nodes valued from [1, N], where node 1 is the root, the task is to count the pairs of nodes having minimum distance between them equal to the difference between the distances of both the nodes from the root.

Examples:

Input: N = 3, Edges[][] = {{1, 2}, {1, 3}}, Below is the Tree:
                        1
                      / \
                   2   3
Output: 5   
Explanation: The following pairs satisfy the conditions: {(1, 1), (2, 2), (3, 3), (2, 1), (3, 1)}.

Input: N = 4, Edges[][] = {{1, 2}, {1, 3}, {2, 4}}, Below is the Tree:
                     1
                   / \
                2   3
               |
             4
Output: 8 

Naive Approach: The simplest approach is to find the distance between all possible pairs (i, j) and check for each pair of nodes (i, j), whether distance(i, j) = distance(1, i) - distance(1, j) or not by using Depth-First Search Traversal. 
Time Complexity: O(N²)
Auxiliary Space: O(N)

Efficient Approach: The above approach can be optimized based on the following observations:

The minimum distance from node X and the ancestor of node X satisfy the above criteria. Therefore, the task is reduced to finding all the ancestors of every node of the tree.

Follow the steps below to solve the problem:

• Initialize a variable, say ans, to store the count of pairs of nodes.
• Perform a DFS Traversal from the root node with arguments current node, parent node, and depth of node up until the current node.
• Recursively perform DFS Traversal on the child node of every node with the current node as parent node and depth increased by 1.
• In every call, add the depth of the current node to the variable ans.
• After completing the above steps, print the value of ans as the result.

Below is the implementation of the above approach:

// C++ program for the above approach

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

// Stores the count of pairs
long long ans = 0;

// Store the adjacency list of
// the connecting vertex
vector<int> adj[int(1e5) + 1];

// Function for theto perform DFS
// traversal of the given tree
void dfsUtil(int u, int par, int depth)
{
    // Traverse the adjacency list
    // of the current node u
    for (auto it : adj[u]) {

        // If the current node
        // is the parent node
        if (it != par) {
            dfsUtil(it, u, depth + 1);
        }
    }

    // Add number of ancestors, which
    // is same as depth of the node
    ans += depth;
}

// Function for DFS traversal of
// the given tree
void dfs(int u, int par, int depth)
{
    dfsUtil(u, par, depth);

    // Print the result
    cout << ans << endl;
}

// Function to find the count of pairs
// such that the minimum distance
// between them is equal to the difference
// between distance of the nodes from root node
void countPairs(vector<vector<int> > edges)
{
    for (int i = 0; i < edges.size(); i++) {

        int u = edges[i][0];
        int v = edges[i][1];

        // Add edges to adj[]
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    dfs(1, 1, 1);
}

// Driver Code
int main()
{
    vector<vector<int> > edges
        = { { 1, 2 }, { 1, 3 }, { 2, 4 } };

    countPairs(edges);

    return 0;
}

// Java program for the above approach
import java.util.*;
class GFG
{

// Stores the count of pairs
static int ans = 0;

// Store the adjacency list of
// the connecting vertex
static ArrayList<Integer> []adj = new ArrayList[(int)(1e5) + 1];
static{
    for(int i = 0; i < adj.length; i++)
        adj[i] = new ArrayList<>();
}
  
// Function for theto perform DFS
// traversal of the given tree
static void dfsUtil(int u, int par, int depth)
{
  
    // Traverse the adjacency list
    // of the current node u
    for (int it : adj[u]) {

        // If the current node
        // is the parent node
        if (it != par) {
            dfsUtil(it, u, depth + 1);
        }
    }

    // Add number of ancestors, which
    // is same as depth of the node
    ans += depth;
}

// Function for DFS traversal of
// the given tree
static void dfs(int u, int par, int depth)
{
    dfsUtil(u, par, depth);

    // Print the result
    System.out.print(ans +"\n");
}

// Function to find the count of pairs
// such that the minimum distance
// between them is equal to the difference
// between distance of the nodes from root node
static void countPairs(int [][]edges)
{
    for (int i = 0; i < edges.length; i++) {

        int u = edges[i][0];
        int v = edges[i][1];

        // Add edges to adj[]
        adj[u].add(v);
        adj[v].add(u);
    }

    dfs(1, 1, 1);
}

// Driver Code
public static void main(String[] args)
{
    int [][]edges
        = { { 1, 2 }, { 1, 3 }, { 2, 4 } };

    countPairs(edges);
}
}

// This code is contributed by 29AjayKumar 

# Python3 program for the above approach

# Stores the count of pairs
ans = 0

# Store the adjacency list of
# the connecting vertex
adj = [[] for i in range(10**5+1)]

# Function for theto perform DFS
# traversal of the given tree
def dfsUtil(u, par, depth):
    global adj, ans
    
    # Traverse the adjacency list
    # of the current node u
    for it in adj[u]:

        # If the current node
        # is the parent node
        if (it != par):
            dfsUtil(it, u, depth + 1)

    # Add number of ancestors, which
    # is same as depth of the node
    ans += depth

# Function for DFS traversal of
# the given tree
def dfs(u, par, depth):
    global ans
    dfsUtil(u, par, depth)
    
    # Print result
    print (ans)

# Function to find the count of pairs
# such that the minimum distance
# between them is equal to the difference
# between distance of the nodes from root node
def countPairs(edges):
    global adj
    for i in range(len(edges)):
        u = edges[i][0]
        v = edges[i][1]

        # Add edges to adj[]
        adj[u].append(v)
        adj[v].append(u)
    dfs(1, 1, 1)

# Driver Code
if __name__ == '__main__':
    edges = [ [ 1, 2 ], [ 1, 3 ], [ 2, 4 ] ]

    countPairs(edges)

    # This code is contributed by mohit kumar 29.

// C# program for the above approach
using System;
using System.Collections.Generic;
public class GFG
{

// Stores the count of pairs
static int ans = 0;

// Store the adjacency list of
// the connecting vertex
static List<int> []adj = new List<int>[(int)(1e5) + 1];
  
// Function for theto perform DFS
// traversal of the given tree
static void dfsUtil(int u, int par, int depth)
{
  
    // Traverse the adjacency list
    // of the current node u
    foreach (int it in adj[u]) {

        // If the current node
        // is the parent node
        if (it != par) {
            dfsUtil(it, u, depth + 1);
        }
    }

    // Add number of ancestors, which
    // is same as depth of the node
    ans += depth;
}

// Function for DFS traversal of
// the given tree
static void dfs(int u, int par, int depth)
{
    dfsUtil(u, par, depth);

    // Print the result
    Console.Write(ans +"\n");
}

// Function to find the count of pairs
// such that the minimum distance
// between them is equal to the difference
// between distance of the nodes from root node
static void countPairs(int [,]edges)
{
    for (int i = 0; i < edges.GetLength(0); i++) 
    {

        int u = edges[i,0];
        int v = edges[i,1];

        // Add edges to []adj
        adj[u].Add(v);
        adj[v].Add(u);
    }

    dfs(1, 1, 1);
}

// Driver Code
public static void Main(String[] args)
{
    int [,]edges
        = { { 1, 2 }, { 1, 3 }, { 2, 4 } };

    for(int i = 0; i < adj.GetLength(0); i++)
        adj[i] = new List<int>();
    countPairs(edges);
}
}

// This code is contributed by 29AjayKumar 

<script>

// Javascript program for the above approach

// Stores the count of pairs
var ans = 0;

// Store the adjacency list of
// the connecting vertex
var adj = Array.from(Array(100001), ()=>Array());

// Function for theto perform DFS
// traversal of the given tree
function dfsUtil(u, par, depth)
{
    // Traverse the adjacency list
    // of the current node u
    adj[u].forEach(it => {
        

        // If the current node
        // is the parent node
        if (it != par) {
            dfsUtil(it, u, depth + 1);
        }
    });

    // Add number of ancestors, which
    // is same as depth of the node
    ans += depth;
}

// Function for DFS traversal of
// the given tree
function dfs(u, par, depth)
{
    dfsUtil(u, par, depth);

    // Print the result
    document.write( ans + "<br>");
}

// Function to find the count of pairs
// such that the minimum distance
// between them is equal to the difference
// between distance of the nodes from root node
function countPairs(edges)
{
    for (var i = 0; i < edges.length; i++) {

        var u = edges[i][0];
        var v = edges[i][1];

        // Add edges to adj[]
        adj[u].push(v);
        adj[v].push(u);
    }

    dfs(1, 1, 1);
}

// Driver Code
var edges
    = [ [ 1, 2 ], [ 1, 3 ], [ 2, 4 ] ];
countPairs(edges);

// This code is contributed by itsok.
</script> 

8

Time Complexity: O(N)
Auxiliary Space: O(N)