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.
Załóżmy, że mysz porusza się ustaloną ścieżką $v_1, \ldots, 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 $dp[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} dp[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(nk)$; 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);
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 $dp2[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 $Dp$ i $Dp2$ spośród wierzchołków na lewo od $u$. Złożoność będzie wtedy $O(nk)$:
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; }