Kefa və park (Həll təhlili)

Kəfa'nın evi olan $1$ təpəsindən dərinlik üzrə axtarışa başlaq. dfs funksiyasının parametrlərindən birində ardıcıl gələn pişiklərlə təpələrin sayını saxlayacağıq. Əgər $v$ təpəsində bu sayı $m$-i keçərsə, bu təpədən dərinlik üzrə axtarışı davam etdirmirik. Ziyarət edilən yarpaqların sayını hesablayırıq. Bir təpə yarpaqdır ki, onun dərəcəsi $1$-ə bərabərdir və özü kök deyil (yəni $1$ təpəsi deyil).

Misal

Birinci və ikinci misallardakı ağacları nəzərdən keçirək.

Birinci misalda Kəfa yalnız bir pişiklə olan təpədən ardıcıl keçə bilər. O, $6$ və $7$ təpələrində yerləşən restoranlara çatacaq.

İkinci misalda Kəfa yalnız iki pişiklə olan təpələrdən ardıcıl keçə bilər. O, $4$ və $5$ təpələrində yerləşən restoranlara çatacaq.

Alqoritmin reallaşdırılması

Pişiklərin yerləşməsini $cats$ massivində saxlayırıq. Ağacı qarşılıqlı əlaqə siyahısında $g$ saxlayırıq.

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

dfs funksiyası $v$ təpəsindən dərinlik üzrə axtarış yerinə yetirir və bu təpədən hansı restoranlara gedilə biləcəyini qaytarır (şərt o ki, bu yolda $m$-dən çox ardıcıl pişiklə təpələr olmasın). $v$ təpəsinin atası $prev$ təpəsidir. $v$ təpəsinə qədər keçilmiş ardıcıl pişiklərlə təpələrin sayı $cnt$-yə bərabərdir.

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

Əgər artıq ardıcıl $m$-dən çox pişiklə keçilibsə, $v$ təpəsindən dərinlik üzrə axtarış davam olunmur.

  if (cnt > m) return 0;

Əgər yarpaqdayıqsa, bilinən restoranların sayını artırırıq.

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

Promennaya $res$, $v$ təpəsindən mümkün olan restoranların sayını hesablayır.

  int res = 0;

$v$ təpəsindən çıxan bütün kənarları (v, to) keçiririk. Əgər $to$ təpəsi $v$ üçün atadırsa, bu təpəyə keçid etmə.

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

Əgər $v$ təpəsində pişik yoxdursa, c = 0 təyin edirik. Əks halda c = cnt + 1. Promennaya $c$ $to$ təpəsinə qədər pişiklə olan ardıcıl təpələrin sayını saxlayır.

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

$to$ təpəsindən dərinlik üzrə axtarışa başlayırıq.

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

Nəticəni qaytarırıq.

  return res;
}

Proqramın əsas hissəsi. Giriş məlumatlarını oxuyarıq.

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

Ağacı qururuq.

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

Cavabı çıxışa veririk. $1$ təpəsindən çağırılan dfs funksiyası məsələyə cavab verir. $1$ təpəsinin atası olmadığı üçün ikinci parametrin qiyməti $-1$ olur. $1$ təpəsindən başlayaraq biz $cats[1]$ pişikləri ziyarət etdik.

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

Python realizasiyası

dfs funksiyası $v$ təpəsindən dərinlik üzrə axtarış yerinə yetirir və bu təpədən hansı restoranlara gedilə biləcəyini qaytarır (şərt o ki, bu yolda $m$-dən çox ardıcıl pişiklə təpələr olmasın). $v$ təpəsinin atası $prev$ təpəsidir. $v$ təpəsinə qədər keçilmiş ardıcıl pişiklərlə təpələrin sayı $cnt$-yə bərabərdir.

def dfs(v, prev, cnt):

Əgər artıq ardıcıl $m$-dən çox pişiklə keçilibsə, $v$ təpəsindən dərinlik üzrə axtarış davam olunmur.

  if cnt > m:
    return 0

Əgər yarpaqdayıqsa, bilinən restoranların sayını artırırıq.

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

Promennaya $res$, $v$ təpəsindən mümkün olan restoranların sayını hesablayır.

  res = 0

$v$ təpəsindən çıxan bütün kənarları (v, to) keçiririk. Əgər $to$ təpəsi $v$ üçün atadırsa, bu təpəyə keçid etmə.

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

Əgər $v$ təpəsində pişik yoxdursa, c = 0 təyin edirik. Əks halda c = cnt + 1. Promennaya $c$ $to$ təpəsinə qədər pişiklə olan ardıcıl təpələrin sayını saxlayır.

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

$to$ təpəsindən dərinlik üzrə axtarışa başlayırıq.

      res += dfs(to, v, c)

Nəticəni qaytarırıq.

  return res

Proqramın əsas hissəsi. Giriş məlumatlarını oxuyarıq.

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

Ağacı qururuq.

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)

Cavabı çıxışa veririk. $1$ təpəsindən çağırılan dfs funksiyası məsələyə cavab verir. $1$ təpəsinin atası olmadığı üçün ikinci parametrin qiyməti $-1$ olur. $1$ təpəsindən başlayaraq biz $cats[1]$ pişikləri ziyarət etdik.

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