Count nodes within K | C ++ 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 ; push node u in queue and initialize its distance as 0 ; loop untill all nodes are processed ; if node is marked , update lastMarked variable ; loop over all neighbors of u and update their distance before pushing in queue ; if not given value already ; return last updated marked value ; method returns count of nodes which are in K - distance range from marked nodes ; vertices in a tree are one more than number of edges ; fill vector for graph ; fill boolean array mark from marked array ; vectors to store distances ; first bfs ( from any random node ) to get one distant marked node ; second bfs to get other distant marked node and also dl is filled with distances from first chosen marked node ; third bfs to fill dr by distances from second chosen marked node ; loop over all nodes ; increase res by 1 , if current node has distance less than K from both extreme nodes ; Driver code to test above methods
#include <bits/stdc++.h> NEW_LINE using namespace std ; int bfsWithDistance ( vector < int > g [ ] , bool mark [ ] , int u , vector < int > & dis ) { int lastMarked ; queue < int > q ; q . push ( u ) ; dis [ u ] = 0 ; while ( ! q . empty ( ) ) { u = q . front ( ) ; q . pop ( ) ; if ( mark [ u ] ) lastMarked = u ; for ( int i = 0 ; i < g [ u ] . size ( ) ; i ++ ) { int v = g [ u ] [ i ] ; if ( dis [ v ] == -1 ) { dis [ v ] = dis [ u ] + 1 ; q . push ( v ) ; } } } return lastMarked ; } int nodesKDistanceFromMarked ( int edges [ ] [ 2 ] , int V , int marked [ ] , int N , int K ) { V = V + 1 ; vector < int > g [ V ] ; 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 ) ; } bool mark [ V ] = { false } ; for ( int i = 0 ; i < N ; i ++ ) mark [ marked [ i ] ] = true ; vector < int > tmp ( V , -1 ) , dl ( V , -1 ) , dr ( V , -1 ) ; u = bfsWithDistance ( g , mark , 0 , tmp ) ; v = bfsWithDistance ( g , mark , u , dl ) ; bfsWithDistance ( g , mark , v , dr ) ; int res = 0 ; for ( int i = 0 ; i < V ; i ++ ) { if ( dl [ i ] <= K && dr [ i ] <= K ) res ++ ; } return res ; } 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 ; }