An easy DFS solution

interview-questions
Tags: #<Tag:0x00007f1817e35090>

#1

vector ans;
vector tg;
vector<vector > g;
void dfs(int s, int par,vector &A){
if(g[s].size()==1 && g[s][0]==par){
ans.push_back(tg[s]);
}
for(auto i: g[s]){
if(i==par) continue;
if(A[i]) tg[i]=tg[s]+1;
else tg[i]=tg[s];
dfs(i,s,A);
}
}

int Solution::solve(vector &A, vector<vector > &B, int C) {
ans.clear();
tg.clear();
g.clear();
int n=A.size();
g.assign(n,vector());
for(int i=0; i<B.size(); i++){
g[B[i][0]-1].push_back(B[i][1]-1);
g[B[i][1]-1].push_back(B[i][0]-1);
}
tg.assign(n,0);
if(A[0]) tg[0]=1;
dfs(0,-1,A);
int val=0;
for(auto i: ans){
if(i<=C) val++;
}
return val;
}