Problem: gold_mine_ch2 (2021 quals)

**Note: This problem shares similarities with [Chapter 1](https://www.facebook.com/codingcompetitions/hacker-cup/2021/qualification-round/problems/C1). The solution to either chapter may help with solving the other, so please consider reading both first.**

Just like [logging](https://www.youtube.com/watch?v=7AheBqBk0iQ), mining is a key industry in Canada's primary sector. Prospectors use geological tools to find ore deposits, then [drill, and blast](https://www.youtube.com/watch?v=zPbY--IXGg8&t=463s) to excavate a system of underground tunnels for access to the ores.

Minerva is a contract miner who spends her days extracting gold ore from underground mines across Ontario. One particular gold mine she's contracting for consists of \(N\) underground caves, numbered from \(1\) to \(N\), with cave \(1\) having an exit to the surface. There are also \(N - 1\) tunnels running among the caves, the \(i\)th of which allows one to travel in either direction between caves \(A_i\) and \(B_i\). Each cave is reachable from cave \(1\) by a sequence of these tunnels.

Cave \(i\) is known to contain \(C_i\) ounces of gold ore. With her trusty dump truck, Minerva plans to drive underground starting in cave \(1\), collect as much gold ore as possible, and exit back to the surface through cave \(1\).

There's just one issue: the tunnels are highly unstable, and will cave in after Minerva drives through them. Therefore, she is only able to traverse each tunnel at most once. Fortunately, the caves themselves are reinforced with pillars and will never collapse — she may visit each cave any number of times (though she can only collect its gold ore at most once).

This safety measure may limit mobility. Thankfully, Minerva's truck is also equipped to be a drill rig. As she travels through the system of caves, she may at any point use it to travel from her current cave to *any* other cave by drilling a new tunnel. **In this chapter, the rig only has the energy to drill at most \(K\) such tunnels.** Unfortunately, these new tunnels will be just as unstable as the existing ones — Minerva will only dare travel through each new tunnel once.

Please help Minerva determine the maximum weight of gold ore that can be collected if she begins in cave \(1\), spends any amount of time exploring the mine (traveling through tunnels, drilling at most \(K\) new tunnels, and collecting ore from any visited caves), and end up once again in cave \(1\).

# Constraints

\(1 \le T \le 30\)
\(1 \le N \le 50\)
\(0 \le K \le 50\)
\(0 \le C_i \le 20{,}000{,}000\)
\(1 \le A_i, B_i \le N\)

The sum of \(N\) across all mines is at most \(1{,}000\).


# Input

Input begins with an integer \(T\), the number of mines that Minerva is contracted to work. For each mine, there is first a line containing the space-separated integers \(N\) and \(K\). Then, there is a second line containing \(N\) space-separated integers, \(C_{1..N}\). Then, \(N - 1\) lines follow, the \(i\)th of which contains the space-separated integers \(A_i\) and \(B_i\).


# Output

For the \(i\)th mine, output a line containing *"Case #i: "* followed by a single integer, the maximum weight of gold ore (in ounces) that Minerva can collect before leaving the mine.


# Sample Explanation

The first mine looks as follows:

{{PHOTO_ID:375673720664225}}

It's possible to visit both caves (collecting \(30\) ounces of ore) and then return to cave \(1\), for example by taking the tunnel from cave \(1\) to cave \(2\) and then drilling a tunnel back to cave \(1\).

In the second mine, the system of caves is the same, but with no ability to drill, Minerva cannot afford to travel to cave \(2\) (as she would become stuck there) and must settle for only collecting \(10\) ounces of ore from cave \(1\).

The third mine looks as follows:

{{PHOTO_ID:541711107152146}}

It's possible to end up with \(3\) ounces of ore, but Minerva cannot collect all \(4\) ounces while still returning to cave \(1\) and only drilling one tunnel.

In the fourth mine, the system of caves is the same, but with the ability to drill twice, Minerva can collect ore from all \(4\) caves. For example, she can take the tunnel to cave \(3\), drill to travel to cave \(2\), drill again to travel to cave \(4\), and finally take the tunnel back to cave \(1\).


Sample input:
7
2 1
10 20
1 2
2 0
10 20
1 2
4 1
1 1 1 1
2 1
4 1
1 3
4 2
1 1 1 1
2 1
4 1
1 3
6 2
5 4 1 3 2 4
5 1
5 4
5 3
5 2
6 3
9 2
2 14 7 6 11 3 6 1 8
4 5
6 7
8 9
1 3
6 8
2 4
4 1
1 8
9 50
9 14 7 6 11 3 6 1 8
4 5
6 7
8 9
1 3
6 8
2 4
4 1
1 8


Sample output:
Case #1: 30
Case #2: 10
Case #3: 3
Case #4: 4
Case #5: 19
Case #6: 50
Case #7: 65


Solution: As in Chapter 1, the mine can be thought of as a rooted tree, with one node per cave, the root being node \(1\), and each tunnel corresponding to an edge.

If \(K = 0\), the answer is \(C_1\). Otherwise, we can think of Minerva's route as consisting of any set of \(K\) edge-disjoint paths in the tree (that is, with no edge shared between multiple paths), with at least one of them including node \(1\). A path including node \(1\) will implicitly correspond to both the first and last portions of her route: we can imagine her starting at node \(1\) and driving to one of the path's endpoints before drilling her first tunnel, and then eventually drilling her final tunnel to the path's other endpoint and driving back to node \(1\). The remaining \(K-1\) paths more directly correspond to the intermediate portions of her route, with her repeatedly drilling to a path's endpoint and then driving along that complete path before drilling another tunnel to the next path. The total value of such a set of paths is equal to the sum of \(C\) values across all nodes which are part of at least one path, and we want to find a set of paths which maximizes this total.

We'll take a dynamic programming (DP) approach. Let \(F(i, j, k)\) be the maximum total value of paths if only node \(i\)'s subtree is considered, such that:
- There are \(j\) paths contained entirely within that subtree
- There's a path with an endpoint at node \(i\) (potentially connecting to \(i\)'s parent) available to be extended further if and only if \(k = 1\) (defined for \(0 \le k \le 1\))

We'll additionally enforce that, if \(i = 1\), at least one of the paths must include node \(i\). Our final answer will be the maximum of \(C_1\) and of \(F(1, j, 0)\) (for \(1 \le j \le K\)).

Computing \(F(i, j, k)\) will require processing \(i\) and its children in a secondary level of DP. Let \(G(i, f, j, k, c)\) be the maximum total value of paths if only node \(i\) and its first \(f\) children are considered, such that:
- There are \(j\) paths contained entirely within those nodes
- There's a path with an endpoint at node \(i\) (potentially connecting to \(i\)'s parent) available to be extended further if and only if \(k = 1\) (defined for \(0 \le k \le 1\))
- There's at least one path connecting \(i\) to one of those \(f\) children if and only if \(c = 1\) (defined for \(0 \le c \le 1\))

We'll recurse through the tree, computing these values. Once we've computed the \(F\) values for all children of a node \(i\), we'll be able to compute node \(i\)'s \(G\) values. After starting with the base case of \(G(i, 0, 0, 0, 0) = 0\), we can compute \(G(f, ...)\) for each \(f\)th child of node \(i\) by combining \(G(f-1, ...)\) with that child's \(F\) values, while considering the options of either having or not having a path connecting node \(i\) and that child. If node \(i\) has \(Z_i\) children, we can then transform \(G(i, Z_i, ...)\) into \(F(i, ...)\). Further details around the DP transitions involved can be found in the official solution's source code.

Overall, it takes a total of \(O(NK^2)\) time to compute all \(O(NK)\) relevant DP values.

[See David Harmeyer's solution video here](https://youtu.be/V1p4ZX-uDM8).


Code: #include <algorithm>
#include <cstring>
#include <iostream>
#include <vector>
using namespace std;

const int LIM = 52;

inline int& setmax(int &l, int r) {
  return l = max(l, r);
}

int N, K, C[LIM];
vector<int> adj[LIM];

// dp[i][j][k] = max value in i's subtree,
// with j new paths present,
// and with a free path ongoing from i's parent if k=1.
int dp[LIM][LIM][2];

// dp2[i][j][k][c] = max value after first i children of current node,
// with j new paths present,
// and with a free path available for use if k=1,
// and with at least one child connected if c=1.
int dp2[LIM][LIM][2][2];

void rec(int i, int parent) {
  int nc = 0;
  for (int j : adj[i]) {
    if (j != parent) {
      rec(j, i);
      nc++;
    }
  }
  // Sub-DP.
  for (int j = 0; j <= nc; j++) {
    memset(dp2[j], -1, sizeof dp2[j]);
  }
  dp2[0][0][0][0] = 0;
  nc = 0;
  for (int j : adj[i]) {
    if (j == parent) {
      continue;
    }
    for (int a1 = 0; a1 <= K; a1++) {
      for (int x : {0, 1}) {
        for (int y : {0, 1}) {
          int d = dp2[nc][a1][x][y];
          if (d < 0) {
            continue;
          }
          for (int a2 = 0; a2 <= K - a1; a2++) {
            // Connect to child.
            int d2 = dp[j][a2][1];
            if (d2 >= 0) {
              setmax(dp2[nc + 1][a1 + a2 + (x ? 0 : 1)][1 - x][1], d + d2);
            }
            // Don't connect to child.
            d2 = dp[j][a2][0];
            if (d2 >= 0) {
              setmax(dp2[nc + 1][a1 + a2][x][y], d + d2);
            }
          }
        }
      }
    }
    nc++;
  }
  // Combine into main DP.
  memset(dp[i], -1, sizeof dp[i]);
  for (int o : {0, 1}) {
    for (int j = 0; j <= K; j++) {
      for (int x : {0, 1}) {
        for (int y : {0, 1}) {
          if (!i && !y) {
            continue;  // Root must connect to at least one child.
          }
          int d = dp2[nc][j][x][y];
          if (d < 0) {
            continue;
          }
          // Include current node's value?
          if (o || y) {
            d += C[i];
          }
          // Free path ongoing from parent?
          setmax(dp[i][j - (o && x ? 1 : 0)][o], d);
        }
      }
    }
  }
}

int solve() {
  cin >> N >> K;
  for (int i = 0; i < N; i++) {
    cin >> C[i];
    adj[i].clear();
  }
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  // DP
  rec(0, -1);
  int ans = C[0];
  for (int i = 0; i <= K; i++) {
    setmax(ans, dp[0][i][0]);
  }
  return ans;
}

int main() {
  int T;
  cin >> T;
  for (int t = 1; t <= T; t++) {
    cout << "Case #" << t << ": " << solve() << endl;
  }
  return 0;
}