Розбір
Start a depth first search from vertex , where Kefa's house is located. One of the parameters of the dfs function will store the number of consecutive vertices with cats. If at vertex this number exceeds , we'll not continue the depth-first search from this vertex. We count the number of visited leaf vertices. A vertex is a leaf if its degree is and it is not a root (vertex ).
Example
Let’s consider the trees from the first and second examples.
In the first example, Kefa can pass through only one consecutive vertex with a cat. He will be able to reach the restaurants located at vertices and .
In the second example, Kefa can pass through two consecutive vertices with cats. He will be able to reach the restaurants located at vertices and .
Algorithm realization
The location of the cats is stored in the array . The tree is stored in the adjacency list .
vector<int> cats; vector<vector<int> > g;
The dfs function performs a depth-first search from vertex and returns the number of restaurants that can be reached from this vertex (under the condition that no more than consecutive vertices with cats are encountered on the path to them). The parent of vertex is the vertex . The number of consecutive vertices with cats encountered up to and including vertex is .
int dfs(int v, int prev, int cnt) {
If more than consecutive cats are already encountered, further depth-first search from vertex is not continued.
if (cnt > m) return 0;
If we are at a leaf, increase the number of accessible restaurants.
if (g[v].size() == 1 && prev != -1) return 1;
In the variable count the number of restaurants reachable from vertex .
int res = 0;
Iterate over all edges originating from vertex . If the vertex is a parent of , do not move to this vertex.
for (int to : g[v]) if (to != prev) {
If there is no cat at the vertex , set . Otherwise, assign . The variable stores the number of consecutive vertices with cats up to and including vertex .
int c = (cats[to] == 0) ? 0 : cnt + 1;
Start a depth-first search from vertex .
res += dfs(to, v, c); }
Return the result.
return res; }
The main part of the program. Read the input data.
scanf("%d %d", &n, &m); cats.resize(n + 1); for (i = 1; i <= n; i++) scanf("%d", &cats[i]);
Construct a tree.
g.resize(n + 1); for (i = 1; i < n; i++) { scanf("%d %d", &a, &b); g[a].push_back(b); g[b].push_back(a); }
Print the result. The dfs function, called from vertex , returns the answer to the problem. Since vertex has no parent, we pass the value as the second parameter. Starting from vertex , we have encountered cats.
printf("%d\n", dfs(1, -1, cats[1]));
Python realization
The dfs function performs a depth-first search from vertex and returns the number of restaurants that can be reached from this vertex (under the condition that no more than consecutive vertices with cats are encountered on the path to them). The parent of vertex is the vertex . The number of consecutive vertices with cats encountered up to and including vertex is .
def dfs(v, prev, cnt):
If more than consecutive cats are already encountered, further depth-first search from vertex is not continued.
if cnt > m: return 0
If we are at a leaf, increase the number of accessible restaurants.
if len(g[v]) == 1 and prev != -1: return 1
In the variable count the number of restaurants reachable from vertex .
res = 0
Iterate over all edges originating from vertex . If the vertex is a parent of , do not move to this vertex.
for to in g[v]: if to != prev:
If there is no cat at the vertex , set . Otherwise, assign . The variable stores the number of consecutive vertices with cats up to and including vertex .
c = 0 if cats[to] == 0 else cnt + 1
Start a depth-first search from vertex .
res += dfs(to, v, c)
Return the result.
return res
The main part of the program. Read the input data.
n, m = map(int, input().split()) cats = [0] + list(map(int, input().split()))
Construct a tree.
g = [[] for _ in range(n + 1)] for _ in range(n - 1): a, b = map(int, input().split()) g[a].append(b) g[b].append(a)
Print the result. The dfs function, called from vertex , returns the answer to the problem. Since vertex has no parent, we pass the value as the second parameter. Starting from vertex , we have encountered cats.
print(dfs(1, -1, cats[1]))