Count nodes within K-distance from all nodes in a set

Given an undirected tree with some marked nodes and a positive number K. We need to print the count of all such nodes which have distance from all marked nodes less than or equal to K that means every node whose distance from all marked nodes is less than or equal to K, should be counted in the result. 

Examples:

In above tree we can see that node with index 
0, 2, 3, 5, 6, 7 have distances less than or equal to 3
from all the marked nodes.
so answer will be 6

We can solve this problem using breadth first search. Main thing to observe in this problem is that if we find two marked nodes which are at largest distance from each other considering all pairs of marked nodes then if a node is at a distance less than K from both of these two nodes then it will be at a distance less than K from all the marked nodes because these two nodes represents the extreme limit of all marked nodes, if a node lies in this limit then it will be at a distance less than K from all marked nodes otherwise not. 
As in above example, node-1 and node-4 are most distant marked node so nodes which are at distance less than 3 from these two nodes will also be at distance less than 3 from node 2 also. 

Now first distant marked node we can get by doing a bfs from any random node, second distant marked node we can get by doing another bfs from marked node we just found from the first bfs and in this bfs we can also found distance of all nodes from first distant marked node and to find distance of all nodes from second distant marked node we will do one more bfs, so after doing these three bfs we can get distance of all nodes from two extreme marked nodes which can be compared with K to know which nodes fall in K-distance range from all marked nodes. 

Implementation:

C++

//  C++ program to count nodes inside K distance
// range from marked nodes
#include <bits/stdc++.h>
using namespace std;

// Utility bfs method to fill distance vector and returns
// most distant marked node from node u
int bfsWithDistance(vector<int> g[], bool mark[], int u,
                                        vector<int>& dis)
{
    int lastMarked;
    queue<int> q;

    //  push node u in queue and initialize its distance as 0
    q.push(u);
    dis[u] = 0;

    //  loop until all nodes are processed
    while (!q.empty())
    {
        u = q.front();      q.pop();
        //  if node is marked, update lastMarked variable
        if (mark[u])
            lastMarked = u;

        // loop over all neighbors of u and update their
        // distance before pushing in queue
        for (int i = 0; i < g[u].size(); i++)
        {
            int v = g[u][i];

            //  if not given value already
            if (dis[v] == -1)
            {
                dis[v] = dis[u] + 1;
                q.push(v);
            }
        }
    }
    //  return last updated marked value
    return lastMarked;
}

// method returns count of nodes which are in K-distance
// range from marked nodes
int nodesKDistanceFromMarked(int edges[][2], int V,
                             int marked[], int N, int K)
{
    //  vertices in a tree are one more than number of edges
    V = V + 1;
    vector<int> g[V];

    //  fill vector for graph
    int u, v;
    for (int i = 0; i < (V - 1); i++)
    {
        u = edges[i][0];
        v = edges[i][1];

        g[u].push_back(v);
        g[v].push_back(u);
    }

    //  fill boolean array mark from marked array
    bool mark[V] = {false};
    for (int i = 0; i < N; i++)
        mark[marked[i]] = true;

    //  vectors to store distances
    vector<int> tmp(V, -1), dl(V, -1), dr(V, -1);

    //  first bfs(from any random node) to get one
    // distant marked node
    u = bfsWithDistance(g, mark, 0, tmp);

    /*  second bfs to get other distant marked node
        and also dl is filled with distances from first
        chosen marked node */
    v = bfsWithDistance(g, mark, u, dl);

    //  third bfs to fill dr by distances from second
    // chosen marked node
    bfsWithDistance(g, mark, v, dr);

    int res = 0;
    //  loop over all nodes
    for (int i = 0; i < V; i++)
    {
        // increase res by 1, if current node has distance
        // less than K from both extreme nodes
        if (dl[i] <= K && dr[i] <= K)
            res++;
    }
    return res;
}

// Driver code to test above methods
int main()
{
    int edges[][2] =
    {
        {1, 0}, {0, 3}, {0, 8}, {2, 3},
        {3, 5}, {3, 6}, {3, 7}, {4, 5},
        {5, 9}
    };
    int V = sizeof(edges) / sizeof(edges[0]);

    int marked[] = {1, 2, 4};
    int N = sizeof(marked) / sizeof(marked[0]);

    int K = 3;
    cout << nodesKDistanceFromMarked(edges, V, marked, N, K);
    return 0;
}

Java

/*package whatever //do not write package name here */
import java.io.*;
import java.util.*;

class GFG
{

  // Java program to count nodes inside K distance
  // range from marked nodes

  // Utility bfs method to fill distance vector and returns
  // most distant marked node from node u
  static int bfsWithDistance(ArrayList<ArrayList<Integer>> g, boolean mark[], int u,
                             ArrayList<Integer> dis)
  {
    int lastMarked = 0;
    Queue<Integer> q = new LinkedList<>();

    // push node u in queue and initialize its distance as 0
    q.add(u);
    dis.set(u , 0);

    // loop until all nodes are processed
    while (!q.isEmpty())
    {
      u = q.remove();
      // if node is marked, update lastMarked variable
      if (mark[u] == true)
        lastMarked = u;

      // loop over all neighbors of u and update their
      // distance before pushing in queue
      for (int i = 0; i < g.get(u).size(); i++)
      {
        int v = g.get(u).get(i);

        // if not given value already
        if (dis.get(v) == -1)
        {
          dis.set(v , dis.get(u) + 1);
          q.add(v);
        }
      }
    }
    // return last updated marked value
    return lastMarked;
  }

  // method returns count of nodes which are in K-distance
  // range from marked nodes
  static int nodesKDistanceFromMarked(int edges[][], int V,
                                      int marked[], int N, int K)
  {
    // vertices in a tree are one more than number of edges
    V = V + 1;
    ArrayList<ArrayList<Integer>>g = new ArrayList<ArrayList<Integer>>(V);
    for(int i=0;i<V;i++){
      g.add(new ArrayList<Integer>());
    }

    // fill vector for graph
    int u, v;
    for (int i = 0; i < (V - 1); i++)
    {
      u = edges[i][0];
      v = edges[i][1];

      g.get(u).add(v);
      g.get(v).add(u);
    }

    // fill boolean array mark from marked array
    boolean mark[] = new boolean[V];
    Arrays.fill(mark, false);
    for (int i = 0; i < N; i++)
      mark[marked[i]] = true;

    // vectors to store distances
    ArrayList<Integer> tmp = new ArrayList<>(),dl = new ArrayList<>(),dr = new ArrayList<>();
    for(int i=0;i<V;i++){
      tmp.add(-1);
      dl.add(-1);
      dr.add(-1);
    }

    // first bfs(from any random node) to get one
    // distant marked node
    u = bfsWithDistance(g, mark, 0, tmp);

    /* second bfs to get other distant marked node
        and also dl is filled with distances from first
        chosen marked node */
    v = bfsWithDistance(g, mark, u, dl);

    // third bfs to fill dr by distances from second
    // chosen marked node
    bfsWithDistance(g, mark, v, dr);

    int res = 0;
    // loop over all nodes
    for (int i = 0; i < V; i++)
    {
      // increase res by 1, if current node has distance
      // less than K from both extreme nodes
      if (dl.get(i) <= K && dr.get(i) <= K)
        res++;
    }
    return res;
  }

  // Driver Code
  public static void main(String args[])
  {
    int edges[][] =
    {
      {1, 0}, {0, 3}, {0, 8}, {2, 3},
      {3, 5}, {3, 6}, {3, 7}, {4, 5},
      {5, 9}
    };
    int V = edges.length;

    int marked[] = {1, 2, 4};
    int N = marked.length;

    int K = 3;
    System.out.println(nodesKDistanceFromMarked(edges, V, marked, N, K));
  }

}

// This code is contributed by shinjanpatra

Python3

# Python3 program to count nodes inside
# K distance range from marked nodes
import queue

# Utility bfs method to fill distance
# vector and returns most distant
# marked node from node u
def bfsWithDistance(g, mark, u, dis):
    lastMarked = 0
    q = queue.Queue()

    # push node u in queue and initialize
    # its distance as 0
    q.put(u)
    dis[u] = 0

    # loop until all nodes are processed
    while (not q.empty()):
        u = q.get()

        # if node is marked, update
        # lastMarked variable
        if (mark[u]):
            lastMarked = u

        # loop over all neighbors of u and
        # update their distance before
        # pushing in queue
        for i in range(len(g[u])):
            v = g[u][i]

            # if not given value already
            if (dis[v] == -1):
                dis[v] = dis[u] + 1
                q.put(v)

    # return last updated marked value
    return lastMarked

# method returns count of nodes which
# are in K-distance range from marked nodes
def nodesKDistanceFromMarked(edges, V, marked, N, K):

    # vertices in a tree are one
    # more than number of edges
    V = V + 1
    g = [[] for i in range(V)]

    # fill vector for graph
    u, v = 0, 0
    for i in range(V - 1):
        u = edges[i][0]
        v = edges[i][1]

        g[u].append(v)
        g[v].append(u)

    # fill boolean array mark from
    # marked array
    mark = [False] * V
    for i in range(N):
        mark[marked[i]] = True

    # vectors to store distances
    tmp = [-1] * V
    dl = [-1] * V
    dr = [-1] * V

    # first bfs(from any random node)
    # to get one distant marked node
    u = bfsWithDistance(g, mark, 0, tmp)

    # second bfs to get other distant
    # marked node and also dl is filled
    # with distances from first chosen
    # marked node
    u = bfsWithDistance(g, mark, u, dl)

    # third bfs to fill dr by distances
    # from second chosen marked node
    bfsWithDistance(g, mark, u, dr)

    res = 0

    # loop over all nodes
    for i in range(V):

        # increase res by 1, if current node
        # has distance less than K from both
        # extreme nodes
        if (dl[i] <= K and dr[i] <= K):
            res += 1
    return res

# Driver Code
if __name__ == '__main__':

    edges = [[1, 0], [0, 3], [0, 8],
             [2, 3], [3, 5], [3, 6],
             [3, 7], [4, 5], [5, 9]]
    V = len(edges)

    marked = [1, 2, 4]
    N = len(marked)

    K = 3
    print(nodesKDistanceFromMarked(edges, V,
                                   marked, N, K))

# This code is contributed by PranchalK

C#

// C# program to count nodes inside K distance
// range from marked nodes

using System;
using System.Collections.Generic;

class GFG {

  // Utility bfs method to fill distance array and returns
  // most distant marked node from node u
  static int BfsWithDistance(List<int>[] g, bool[] mark,
                             int u, int[] dis)
  {
    int lastMarked = 0;
    Queue<int> q = new Queue<int>();

    // push node u in queue and initialize its distance
    // as 0
    q.Enqueue(u);
    dis[u] = 0;

    // loop until all nodes are processed
    while (q.Count > 0) {
      u = q.Dequeue();
      // if node is marked, update lastMarked variable
      if (mark[u])
        lastMarked = u;

      // loop over all neighbors of u and update their
      // distance before pushing in queue
      for (int i = 0; i < g[u].Count; i++) {
        int v = g[u][i];

        // if not given value already
        if (dis[v] == -1) {
          dis[v] = dis[u] + 1;
          q.Enqueue(v);
        }
      }
    }
    // return last updated marked value
    return lastMarked;
  }

  // method returns count of nodes which are in K-distance
  // range from marked nodes
  static int NodesKDistanceFromMarked(int[, ] edges,
                                      int V, int[] marked,
                                      int N, int K)
  {
    // vertices in a tree are one more than number of
    // edges
    V = V + 1;
    List<int>[] g = new List<int>[ V ];
    for (int i = 0; i < V; i++)
      g[i] = new List<int>();

    int u, v;
    for (int i = 0; i < (V - 1); i++) {
      u = edges[i, 0];
      v = edges[i, 1];

      g[u].Add(v);
      g[v].Add(u);
    }

    // fill boolean array mark from marked array
    bool[] mark = new bool[V];
    for (int i = 0; i < N; i++)
      mark[marked[i]] = true;

    // arrays to store distances
    int[] tmp = new int[V];
    int[] dl = new int[V];
    int[] dr = new int[V];
    for (int i = 0; i < V; i++) {
      tmp[i] = dl[i] = dr[i] = -1;
    }

    // first bfs(from any random node) to get one
    // distant marked node
    u = BfsWithDistance(g, mark, 0, tmp);

    /* second bfs to get other distant marked node
            and also dl is filled with distances from first
            chosen marked node */
    v = BfsWithDistance(g, mark, u, dl);

    // third bfs to fill dr by distances from second
    // chosen marked node
    BfsWithDistance(g, mark, v, dr);

    int res = 0;
    // loop over all nodes
    for (int i = 0; i < V; i++) {
      // increase res by 1, if current node has
      // distance less than K from both extreme nodes
      if (dl[i] <= K && dr[i] <= K)
        res++;
    }
    return res;
  }

  // Driver code to test above methods
  static void Main(string[] args)
  {
    int[, ] edges = { { 1, 0 }, { 0, 3 }, { 0, 8 },
                     { 2, 3 }, { 3, 5 }, { 3, 6 },
                     { 3, 7 }, { 4, 5 }, { 5, 9 } };
    int V = edges.GetLength(0);
    int[] marked = { 1, 2, 4 };
    int N = marked.Length;

    int K = 3;
    int ans = NodesKDistanceFromMarked(edges, V, marked,
                                       N, K);
    Console.WriteLine(ans);

  }
}

// This code is contributed by cavi4762.

Javascript

<script>
        /* Javascript program to count nodes inside K distance
        range from marked nodes*/


        // Utility bfs method to fill distance vector and returns
        // most distant marked node from node u
        function bfsWithDistance(g, mark, u, dis)
        {
            let lastMarked = 0;
            let q = [];

            // push node u in queue and initialize its distance as 0
            q.push(u);
            dis[u] = 0;

            // loop until all nodes are processed
            while (q.length > 0)
            {
                u = q.shift()
                // if node is marked, update lastMarked variable
                if (mark[u])
                    lastMarked = u;

                // loop over all neighbors of u and update their
                // distance before pushing in queue
                for (let i = 0; i < g[u].length; i++)
                {
                    let v = g[u][i];

                    // if not given value already
                    if (dis[v] == -1)
                    {
                        dis[v] = dis[u] + 1;
                        q.push(v);
                    }
                }
            }
            // return last updated marked value
            return lastMarked;
        }

        // method returns count of nodes which are in K-distance
        // range from marked nodes
        function nodesKDistanceFromMarked(edges, V, marked, N, K)
        {
            // vertices in a tree are one more than number of edges
            V = V + 1;
            let g = new Array(V);
            for(let i = 0; i < V; i++)
                g[i] = [];

            for (let i = 0; i < (V - 1); i++)
            {
                let u = edges[i][0];
                let v = edges[i][1];

                g[u].push(v);
                g[v].push(u);
            }

            // fill boolean array mark from marked array
            let mark = new Array(V);
            for (let i = 0; i < N; i++)
                mark[marked[i]] = true;

            // vectors to store distances
            let tmp = [], dl = [], dr = [];
            for(let i = 0; i < V; i++)
            {
                tmp[i] = dl[i] = dr[i] = -1;
            }

            // first bfs(from any random node) to get one
            // distant marked node
            u = bfsWithDistance(g, mark, 0, tmp);

            /* second bfs to get other distant marked node
                and also dl is filled with distances from first
                chosen marked node */
            v = bfsWithDistance(g, mark, u, dl);

            // third bfs to fill dr by distances from second
            // chosen marked node
            bfsWithDistance(g, mark, v, dr);

            let res = 0;
            // loop over all nodes
            for (let i = 0; i < V; i++)
            {

                // increase res by 1, if current node has distance
                // less than K from both extreme nodes
                if (dl[i] <= K && dr[i] <= K)
                    res++;
            }
            return res;
        }

        // Driver code to test above methods
        let edges=
        [
            [1, 0], [0, 3], [0, 8], [2, 3],
            [3, 5], [3, 6], [3, 7], [4, 5],
            [5, 9]
        ];
        let V = edges.length;
        let marked = [1, 2, 4];
        let N = marked.length;

        let K = 3;
        let ans = nodesKDistanceFromMarked(edges, V, marked, N, K);
        document.write(ans);

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

Output

6

Time Complexity: O(E+V) (as we do 2 times bfs so overall time complexity is as same as time complexity for BFS)

Space Complexity: O(V) (auxiliary space required for queue todo BFS)