Кефа і парк (Розбір)

Запустим пошук вглиб з вершини $1$, де знаходиться дім Кефи. В одному з параметрів функції dfs будемо зберігати кількість поспіль вершин з котами. Якщо у вершині $v$ це число перевищить $m$, подальший пошук вглиб з цієї вершини не виконуємо. Підраховуємо кількість відвідуваних листків. Вершина є листком, якщо її ступінь дорівнює $1$ і при цьому вона не є коренем (тобто вершиною $1$).

Приклад

Розглянемо дерева з першого та другого прикладів.

У першому прикладі Кефа може пройти поспіль тільки через одну вершину з котом. Він зможе потрапити у ресторани, розташовані у вершинах $6$ та $7$.

У другому прикладі Кефа може пройти поспіль тільки через дві вершини з котами. Він зможе потрапити у ресторани, розташовані у вершинах $4$ та $5$.

Реалізація алгоритму

Розташування котів зберігаємо в масиві $cats$. Дерево зберігаємо у списку суміжності $g$.

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

Функція dfs виконує пошук вглиб з вершини $v$ і повертає кількість ресторанів, до яких можна дістатися з цієї вершини (за умови що на шляху до них зустрічається не більше $m$ поспіль вершин з котами). Попередником вершини $v$ є вершина $prev$. Кількість поспіль вершин з котами, пройдених включно до вершини $v$, дорівнює $cnt$.

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

Якщо вже пройдено більше $m$ котів поспіль, то подальший пошук вглиб з вершини $v$ не продовжується.

  if (cnt > m) return 0;

Якщо ми знаходимося у листі, то збільшуємо кількість ресторанів, до яких можна потрапити.

  if (g[v].size() == 1 && prev != -1) return 1;

У змінній $res$ підраховуємо кількість ресторанів, досяжних з вершини $v$.

  int res = 0;

Перебираємо всі ребра (v, to), що виходять з $v$. Якщо вершина $to$ є попередником $v$, то не переходимо в цю вершину.

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

Якщо у вершині $v$ немає кота, встановлюємо c = 0. У противному випадку присвоюємо c = cnt + 1. Змінна $c$ зберігає кількість поспіль вершин з котами до вершини $to$ включно.

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

Запускаємо пошук вглиб з вершини $to$.

      res += dfs(to, v, c);
    }

Повертаємо результат.

  return res;
}

Основна частина програми. Читаємо вхідні дані.

scanf("%d %d", &n, &m);
cats.resize(n + 1);
for (i = 1; i <= n; i++)
  scanf("%d", &cats[i]);

Будуємо дерево.

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);
}

Виводимо відповідь. Функція dfs, викликана з вершини $1$, повертає відповідь на задачу. Оскільки у вершини $1$ немає предка, другим параметром передаємо значення $-1$. Розпочавши з вершини $1$, ми відвідали $cats[1]$ котів.

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

Python реалізація

Функція dfs виконує пошук вглиб з вершини $v$ і повертає кількість ресторанів, до яких можна дістатися з цієї вершини (за умови що на шляху до них зустрічається не більше $m$ поспіль вершин з котами). Попередником вершини $v$ є вершина $prev$. Кількість поспіль вершин з котами, пройдених включно до вершини $v$, дорівнює $cnt$.

def dfs(v, prev, cnt):

Якщо вже пройдено більше $m$ котів поспіль, то подальший пошук вглиб з вершини $v$ не продовжується.

  if cnt > m:
    return 0

Якщо ми знаходимося у листі, то збільшуємо кількість ресторанів, до яких можна потрапити.

  if len(g[v]) == 1 and prev != -1:
    return 1

У змінній $res$ підраховуємо кількість ресторанів, досяжних з вершини $v$.

  res = 0

Перебираємо всі ребра (v, to), що виходять з $v$. Якщо вершина $to$ є попередником $v$, то не переходимо в цю вершину.

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

Якщо у вершині $v$ немає кота, встановлюємо c = 0. У противному випадку присвоюємо c = cnt + 1. Змінна $c$ зберігає кількість поспіль вершин з котами до вершини $to$ включно.

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

Запускаємо пошук вглиб з вершини $to$.

      res += dfs(to, v, c)

Повертаємо результат.

  return res

Основна частина програми. Читаємо вхідні дані.

n, m = map(int, input().split())
cats = [0] + list(map(int, input().split()))

Будуємо дерево.

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)

Виводимо відповідь. Функція dfs, викликана з вершини $1$, повертає відповідь на задачу. Оскільки у вершини $1$ немає предка, другим параметром передаємо значення $-1$. Розпочавши з вершини $1$, ми відвідали $cats[1]$ котів.

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