Editorial

Start a depth first search from vertex $1$, 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 $v$ this number exceeds $m$, 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 $1$ and it is not a root (vertex $1$).

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 $6$ and $7$.

In the second example, Kefa can pass through two consecutive vertices with cats. He will be able to reach the restaurants located at vertices $4$ and $5$.

Algorithm realization

The location of the cats is stored in the array $cats$. The tree is stored in the adjacency list $g$.

vector<int> cats;
vector<vector<int> > g;

The dfs function performs a depth-first search from vertex $v$ and returns the number of restaurants that can be reached from this vertex (under the condition that no more than $m$ consecutive vertices with cats are encountered on the path to them). The parent of vertex $v$ is the vertex $prev$. The number of consecutive vertices with cats encountered up to and including vertex $v$ is $cnt$.

int dfs(int v, int prev, int cnt)
{

If more than $m$ consecutive cats are already encountered, further depth-first search from vertex $v$ 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 $res$ count the number of restaurants reachable from vertex $v$.

  int res = 0;

Iterate over all edges $(v,to)$ originating from vertex $v$. If the vertex $to$ is a parent of $v$, do not move to this vertex.

  for (int to : g[v])
    if (to != prev) 
    {

If there is no cat at the vertex $v$, set $c=0$. Otherwise, assign $c=cnt+1$. The variable $c$ stores the number of consecutive vertices with cats up to and including vertex $to$.

      int c = (cats[to] == 0) ? 0 : cnt + 1;

Start a depth-first search from vertex $to$.

      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 $1$, returns the answer to the problem. Since vertex $1$ has no parent, we pass the value $-1$ as the second parameter. Starting from vertex $1$, we have encountered $cats[1]$ cats.

printf("%d\n", dfs(1, -1, cats[1]));

Python realization

The dfs function performs a depth-first search from vertex $v$ and returns the number of restaurants that can be reached from this vertex (under the condition that no more than $m$ consecutive vertices with cats are encountered on the path to them). The parent of vertex $v$ is the vertex $prev$. The number of consecutive vertices with cats encountered up to and including vertex $v$ is $cnt$.

def dfs(v, prev, cnt):

If more than $m$ consecutive cats are already encountered, further depth-first search from vertex $v$ 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 $res$ count the number of restaurants reachable from vertex $v$.

  res = 0

Iterate over all edges $(v,to)$ originating from vertex $v$. If the vertex $to$ is a parent of $v$, do not move to this vertex.

  for to in g[v]:
    if to != prev:

If there is no cat at the vertex $v$, set $c=0$. Otherwise, assign $c=cnt+1$. The variable $c$ stores the number of consecutive vertices with cats up to and including vertex $to$.

      c = 0 if cats[to] == 0 else cnt + 1

Start a depth-first search from vertex $to$.

      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 $1$, returns the answer to the problem. Since vertex $1$ has no parent, we pass the value $-1$ as the second parameter. Starting from vertex $1$, we have encountered $cats[1]$ cats.

print(dfs(1, -1, cats[1]))