Mysz zaczyna w dowolnym wierzchołku, następnie przechodzi pewną ścieżką i opuszcza drzewo w dowolnym wierzchołku. Może przy tym $k$ razy „zaznaczyć” wierzchołek.

Gdy wchodzi do wierzchołka to „spotyka” znajdujące się tam gołębie. Gdy zaznacza wierzchołek $v$ , to wszystkie gołębie z sąsiednich wierzchołków przelatują na wierzczhołek $v$ .

Następnie kot przechodzi po tej samej ścieżce.

Naszym zadaniem jest tak wybrać ścieżkę i zaznaczanie wierzchołków, żeby zmaksymalizować wartość: liczba gołębi spotkanych przez kota minus liczba gołębi spotkanych przez mysz.

Początkowe obserwacje

Załóżmy, że mysz porusza się ustaloną ścieżką $v_{1}, \dots, v_{s}$ i zastanówmy się, jak ma zaznaczać wierzchołki. Interesuje nas jedynie liczba gołębi w wierzchołkach ścieżki: $p [v_{i}]$ oraz liczba gołębi w sąsiadach wierzchołka $v_{i}$ , które nie leżą na ścieżce; oznaczmy to przez $q [v_{i}]$ .

Zauważmy, że mysz będzie spotykać tylko gołębie ze ścieżki (choć być może nie wszystkie). Natomiast kot spotka wszystkie gołębie ze ścieżki i być może niektóre z wierzchołków sąsiednich.

Jeżeli mysz jest w wierzchołku $v_{i}$ , to spotyka tam $p [v_{i}]$ gołębi, jeśli nie zaznaczyła wierzchołka $v_{i - 1}$ , albo 0 gołębi, jeśli zaznaczyła $v_{i - 1}$ .

Z kolei dla kota istotne są gołębie z sąsiadów: jeśli mysz zaznaczy $v_{i}$ , to kot spotka $q [v_{i}]$ dodatkowych gołębi, a jeśli mysz nie zaznaczy $v_{i}$ , to nie.

Podsumowując: zaznaczenie wierzchołka $v_{i}$ spowoduje, że koszt na ścieżce zmieni się o $q [v_{i}] + p [v_{i + 1}]$ (bo kot spotyka $q [v_{i}]$ gołębi oraz mysz nie spotyka $p [v_{i + 1}]$ gołębi, zatem te spotkane przez kota będą się liczyć). Niezaznaczenie $v_{i}$ nie zmieni kosztu (spotkane $p [v_{i + 1}]$ gołębi przez mysz i przez kota zniosą się wzajemnie).

Można zauważyć, że jeśli ukorzenimy drzewo w wierzchołku $v_{1}$ , to nasza ścieżka będzie prowadzić w dół drzewa. Co więcej, wartość $q [v_{i}] + p [v_{i + 1}]$ jest po prostu sumą liczby gołębi ze wszystkich dzieci wierzchołka $v_{i}$ .

Wystarczy teraz zrobić programowanie dynamiczne o stanie $d p [v] [j]$ , który oznacza koszt, jeśli rozważyliśmy najlepszą ścieżkę z $v$ w dół drzewa i użyliśmy na niej $j$ zaznaczeń wierzchołków. Odpowiedzią jest wtedy $max_{j \leq k} d p [v_{1}] [j]$ :

const int N = 110000, V = 110;
int n,k,p[N];
vector<int> adj[N];
ll dp[N][V];

void dfs(int v, int par) {
  ll sum=0;
  for (int u : adj[v]) if (u != par) {
    dfs(u, v);
    sum += p[u];
  }

  dp[v][0] = 0;
  REP(i,k) {
    dp[v][i+1] = 0;
    for (int u : adj[v]) if (u != par) {
      dp[v][i+1] = max(dp[v][i+1], sum + dp[u][i]);
      dp[v][i+1] = max(dp[v][i+1], dp[u][i+1]);
    }
  }
}

int main() {
  scanf("%d%d",&n,&k);
  REP(i,n) scanf("%d",&p[i]);
  REP(i,n-1) {
    int a,b; scanf("%d%d",&a,&b);
    --a; --b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }

  dfs(0, -1);

  ll ans = 0;
  REP(i,k) {
    ans = max(ans, dp[0][i+1]);
  }
  printf("%lld\n", ans);
}

Złożoność czasowa tego rozwiązania to $O (n k)$ ; zalicza ono podzadanie 3.

Łatwo jest przerobić to rozwiązanie na program o złożoności $O (n^{2} k)$ , który rozwiązuje podzadanie 2. Wystarczy rozważyć wszystkie możliwości wyboru wierzchołka $v_{1}$ :

  ll ans = 0;

  REP(root,n) {
    dfs(root, -1);
    REP(i,k) {
      ans = max(ans, dp[root][i+1]);
    }
  }
  printf("%lld\n", ans);

Rozwiązanie wzorcowe

Jeżeli ścieżka nie musi zaczynać się w korzeniu, to może zaczynać się w pewnym wierzchołku $v$ i iść w dół, albo zaczynać się w pewnym wierzchołku i iść w górę do $v$ , albo zaczynać się w pewnym wierzchołku, iść w górę do wierzchołka $w$ , a potem z niego w dół.

Analogicznie zatem możemy programowaniem dynamicznym wyznaczyć stany $d p 2 [v] [j]$ dla ścieżek idących w górę. Przy czym należy tutaj uważać, bo zaznaczając wierzchołek $v_{i}$ na takiej ścieżce, musimy dodać do wyniku sumę liczby gołębi z dzieci tego wierzchołka oprócz $v_{i - 1}$ plus suma z rodzica.

Aby obliczyć przypadek trzeci (czyli ścieżki przechodzące przez $w$ ), dla każdego $w$ rozpatrujemy jego dwoje dzieci $u_{1}$ i $u_{2}$ , i rozważamy najlepszą ścieżkę w górę do $u_{1}$ oraz ścieżkę z $u_{2}$ w dół, przy czym mamy $O (k)$ możliwości na rozdysponowanie zaznaczanych wierzchołków pomiędzy tymi ścieżkami. Trzeba jeszcze uważnie rozważyć przypadki gdy zaznaczamy lub nie wierzchołek $w$ .

Szczegóły w poniższym kodzie:

void dfs(int v, int par) {
  ll sum=0;
  for (int u : adj[v]) if (u != par) {
    dfs(u, v);
    sum += p[u];
  }
  int ppar = par != -1 ? p[par] : 0;

  dp[v][0] = 0;
  dp2[v][0] = 0;
  REP(i,k) {
    dp[v][i+1] = 0;
    dp2[v][i+1] = ppar;
    for (int u : adj[v]) if (u != par) {
      dp[v][i+1] = max(dp[v][i+1], sum + dp[u][i]);
      dp[v][i+1] = max(dp[v][i+1], dp[u][i+1]);

      dp2[v][i+1] = max(dp2[v][i+1], sum - p[u] + ppar + max(dp2[u][i], (ll)p[u]));
      dp2[v][i+1] = max(dp2[v][i+1], dp2[u][i+1]);
    }
  }

  REP(i,k+1) {
    for (int u1 : adj[v]) if (u1 != par) {
      ans = max(ans, dp[u1][i]);
      if (i) {
        ans = max(ans, dp[u1][i-1] + ppar + sum);
      }
      ans = max(ans, dp2[u1][i]);
      for (int u2 : adj[v]) if (u2 != par && u1 != u2) {
        ans = max(ans, dp[u1][i] + dp2[u2][k-i]);
        if (i) {
          ans = max(ans, dp[u1][i-1] + ppar + sum - p[u2] + dp2[u2][k-i]);
        }
      }
    }
  }
}

Złożoność czasowa tego rozwiązania to $O (n^{2} k)$ z uwagi na kwadratową pętlę przebiegającą po wierzchołkach $u_{1}$ i $u_{2}$ . Ale można to przyspieszyć, iterując się po $u$ i pamiętając w osobnej tablicy najlepsze $D p$ i $D p 2$ spośród wierzchołków na lewo od $u$ . Złożoność będzie wtedy $O (n k)$ :

  REP(i,k+1) {
    Dp[i] = Dp2[i] = 0;
  }

  bool first = 1;
  for (int u : adj[v]) if (u != par) {
    REP(i,k+1) {
      ans = max(ans, dp[u][i]);
      if (i) {
        ans = max(ans, dp[u][i-1] + ppar + sum);
      }
      ans = max(ans, dp2[u][i]);
      if (!first) {
        ans = max(ans, dp[u][i] + Dp2[k-i]);
        ans = max(ans, Dp[i] + dp2[u][k-i]);
        if (i) {
          ans = max(ans, dp[u][i-1] + ppar + sum - p[Dp2_v[k-i]] + Dp2[k-i]);
          ans = max(ans, Dp[i-1] + ppar + sum - p[u] + dp2[u][k-i]);
        }
      }
    }
    REP(i,k+1) {
      Dp[i] = max(Dp[i], dp[u][i]);
      if (dp2[u][i] > Dp2[i]) {
        Dp2[i] = dp2[u][i];
        Dp2_v[i] = u;
      }
    }
    first = 0;
  }