Problem: chainblock (2021 round2)

After running tests on samples of ore, mining engineers must [plot out the economics of extraction, transportation, and treatment in long range terms](https://youtu.be/zPbY--IXGg8?t=618s). Feasibility reports are then drafted for executive decision.

In a given Ontario town, all the extracted ore is shipped using the town's transportation network, known as its Bulk Transportation Chain (BTC). The BTC is a network of \(N\) ore processing plants, numbered from \(1\) to \(N\), with \(N-1\) roads running amongst them. The \(i\)th road allows ore to be shipped in either direction between plants \(A_i\) and \(B_i\). It is possible to reach any plant from any other plant by following a sequence of roads.

Mining is a dirty business, and not just among blue-collar workers. Executives of a new mining project are plotting to sabotage their rival company's presence in a given town by disrupting its transportation operations. However, in trying to actually build the blockade on the aforementioned roads, their plan was foiled.

As it turns out, ore processing plants communicate with one another via radio towers. Plant \(i\) within the BTC network operates their radio towers on a frequency of \(F_i\) kilohertz, where \(F_i\) is an integer. On a given radio frequency, all of the plants will be able to communicate freely. Thus if ore transportation is disrupted between any pair of plants on the same frequency, they will notice immediately and alert their headquarters.

Equipped with this new intel, the executives would like to take another crack at disrupting their rival's operations. This time, they would like to blockade as many roads as they can, while ensuring that any pair of plants on the same radio frequency can still transport ore to each other through some sequence of roads. This way, every single plant will remain ignorant to any blockaded roads.

Given this restriction, what is the maximum number of roads that can be blockaded?

# Constraints

\(1 \le T \le 45\)
\(2 \le N \le 800{,}000\)
\(1 \le A_i, B_i \le N\)
\(1 \le F_i \le N\)

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


# Input

Input begins with an integer \(T\), the number of towns that the execs are scoping out. For each town, there is first a line containing the single integer \(N\). Then, \(N-1\) lines will follow, the \(i\)th of which contains the \(2\) space-separated integers \(A_i\) and \(B_i\). Then, there is a line with \(N\) space-separated integers, \(F_{1..N}\).


# Output

For the \(i\)th town, output a line containing *"Case #i: "* followed by a single integer, the maximum number of roads that can be blockaded without the rival company finding out.


# Sample Explanation

In the first town, the BTC network looks as follows:

{{PHOTO_ID:378165097140012}}

Since both plants are on the same frequency, blockading the road between them would lead to the rival company finding out.

The second town is similar to the first, except the plants are on different frequencies. Therefore, they would not be able to detect the blockade.

In the third town, the BTC network looks as follows:

{{PHOTO_ID:449066296418790}}

The fourth town is similar to the third, except that neither road may be blocked without disconnecting equal-frequency plants \(2\) and \(3\).

In the fifth town, at most \(2\) roads may be blockaded (the roads between plants \(2\) and \(5\), and between plants \(5\) and \(1\)).


Sample input:
6
2
1 2
1 1
2
1 2
1 2
3
1 2
1 3
2 2 3
3
1 2
1 3
2 3 3
6
2 5
4 1
1 3
5 1
3 6
2 3 5 5 6 2
10
10 7
2 8
7 9
7 3
1 3
3 8
8 5
5 6
6 4
5 3 3 4 6 3 10 6 8 8


Sample output:
Case #1: 0
Case #2: 1
Case #3: 1
Case #4: 0
Case #5: 2
Case #6: 3


Solution: The network of plants can be thought of as a tree, which we can arbitrarily root at node (plant) \(1\). We can then observe that an edge (road) between a node \(i\) and its parent node \(P_i\) (if any) may be blockaded if and only if, for each frequency \(f\), either *all* or *none* of the nodes which have a frequency of \(f\) are located within \(i\)'s subtree. Otherwise, there would be at least one pair of nodes with frequency \(f\) (with one inside \(i\)'s subtree and one outside it) which would become inaccessible if the edge in question were blockaded. Note that these condition are independent per edge. What remains is checking these conditions per edge, with one approach for doing so described below.

For a given frequency \(f\), let \(S_f\) be the set of nodes with frequency \(f\). All edges on paths between all pairs of nodes in \(S_f\) cannot be blockaded. Equivalently, letting \(L_f\) be the [lowest common ancestor (LCA)](https://en.wikipedia.org/wiki/Lowest_common_ancestor) of all nodes in \(S_f\), then all edges on paths between \(L_f\) and nodes in \(S_f\) cannot be blockaded. Note that the LCA of \(3\) or more nodes can be computed in terms of LCAs of pairs of nodes – for example, we can compute \(LCA(S_{f,1}, S_{f,2})\), \(LCA(S_{f,2}, S_{f,3})\), and so on, and then take the one with the minimum depth. [Tarjan's off-line LCA algorithm](https://en.wikipedia.org/wiki/Tarjan%27s_off-line_lowest_common_ancestors_algorithm) can be used to compute all of the required LCAs in a total of \(O(N)\) time, though an online \(O(N \log(N))\) algorithm would also suffice.

Let \(D_i\) be the depth of node \(i\), and \(Z_i\) be the minimum depth of \(L_f\) for any frequency \(f\) found within \(i\)'s subtree. The edge between node \(i\) and \(P_i\) (if any) may then be blockaded if and only if \(D_i = Z_i\).

In order to compute all nodes' \(Z\) values efficiently, let's begin by computing \(L_f\) for each relevant frequency \(f\), and for each node \(i\) in \(S_f\), just setting \(Z_i\) to equal \(D_{L_f}\) (even though this may change later). Let's then recurse through the tree and process the nodes in post-order (such that each node \(i\)'s children are processed before \(i\) is processed), such that by the time we process each node \(i\), we expect \(Z_i\) to be accurate. When processing node \(i\), all we need to do is propagate its \(Z\) value up to its parent node \(P_i\), as all frequencies within \(i\)'s subtree are also within \(P_i\)'s subtree. In particular, we should update \(Z_{P_i}\) to equal \(\min(Z_{P_i}, Z_i)\).

This algorithm takes a total of as little as \(O(N)\) time.

[See David Harmeyer (SecondThread)'s solution video here.](https://youtu.be/0SR2rrsJLTM)


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

const int LIM = 800008;
const int LIM2 = 21; // ceil(log2(LIM)) + 1

int N;
int D[LIM], A[LIM];
vector<int> adj[LIM], F[LIM], P[LIM];
vector<int> O;

void rec(int i) {
  O.push_back(i);
  for (int j = 0; j < LIM2; j++) {
    int p = P[i][j];
    if (p < 0) {
      break;
    }
    P[i][j + 1] = P[p][j];
  }
  for (int j : adj[i]) {
    if (j != P[i][0]) {
      P[j][0] = i;
      D[j] = D[i] + 1;
      rec(j);
    }
  }
}

int lca(int a, int b) {
  if (D[a] < D[b]) {
    swap(a, b);
  }
  for (int i = LIM2 - 1; i >= 0; i--) {
    int p = P[a][i];
    if (p >= 0 && D[p] >= D[b]) {
      a = p;
    }
  }
  for (int i = LIM2 - 1; i >= 0; i--) {
    if (P[a][i] != P[b][i]) {
      a = P[a][i];
      b = P[b][i];
    }
  }
  return a == b ? a : P[a][0];
}

int solve() {
  O.clear();
  for (int i = 0; i < LIM; i++) {
    adj[i].clear();
    P[i].assign(LIM2, -1);
    F[i].clear();
  }
  // Input.
  cin >> N;
  for (int i = 0, a, b; i < N - 1; i++) {
    cin >> a >> b;
    a--;
    b--;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  for (int i = 0, f; i < N; i++) {
    cin >> f;
    F[f - 1].push_back(i);
  }
  // DFS through tree, computing depths, ancestors, and pre-order.
  rec(0);
  // Compute LCA requirements for each frequency.
  for (int i = 0; i < N; i++) {
    if (F[i].empty()) {
      continue;
    }
    int a = F[i][0];
    for (int j = 1; j < (int)F[i].size(); j++) {
      a = lca(a, F[i][j]);
    }
    for (int k : F[i]) {
      A[k] = D[a];
    }
  }
  // Iterate upwards through tree, computing answer.
  int ans = 0;
  for (int j = N - 1; j > 0; j--) {
    int i = O[j];
    if (A[i] == D[i]) {
      ans++;
    } else {
      A[P[i][0]] = min(A[P[i][0]], A[i]);
    }
  }
  return ans;
}

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