Problem: blockchain (2021 round1)

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\), and has the capacity for at most \(C_i\) kilograms of ore to be shipped through it per minute. It is possible to reach any plant from any other plant by following a sequence of roads.

The *shipping capacity* for an unordered pair of processing plants is the maximum weight of ore that may be shipped between them (in either direction) per minute. This is equal to the minimum capacity of any road on the shortest sequence of roads connecting them (such a sequence is unique per pair of plants).

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. They hope to blockade a single one of the \(N-1\) roads within the town's BTC network, preventing any ore from being shipped along it and therefore reducing its capacity to \(0\). They would like to negatively impact the shipping capacities between various pairs of processing plants as much as possible, but have yet to decide which road would best accomplish their goal, so they'll consider each of them in turn.

Let \(S_i\) be the sum of shipping capacities, in kilograms, across all \(N*(N-1)/2\) unordered pairs of processing plants if the \(i\)th road were to be blockaded. In order to reduce the size of the output, you must calculate a single integer: the product of \(S_{1..(N-1)}\). As this product may be very large, you should only compute its value modulo \(1{,}000{,}000{,}007\).


# Constraints

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

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 a single integer \(N\). Then, \(N-1\) lines will follow, the \(i\)th of which contains the \(3\) space-separated integers \(A_i\), \(B_i\), and \(C_i\).


# Output

For the \(i\)th town, output a line containing *"Case #i: "* followed by a single integer, the product of \(S_{1..(N-1)}\), modulo \(1{,}000{,}000{,}007\).


# Sample Explanation

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

{{PHOTO_ID:654301598882147}}

If the single road were blockaded (with its capacity reduced to \(0\)), the shipping capacity between plants \(1\) and \(2\) would become 0. Therefore, \(S = [0]\), and the final answer is \(0\).

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

{{PHOTO_ID:4343232682381910}}

If the first road were blockaded, the shipping capacity between plants \(2\) and \(3\) would be \(5\), with the shipping capacities for other pairs of plants (plants \(1\) and \(2\), and plants \(1\) and 3) reduced to \(0\). If the second road were blockaded, there would instead just be a shipping capacity of \(4\) between plants \(1\) and \(2\). Therefore, \(S = [5, 4]\), and the answer is \((5 * 4)\) modulo \(1{,}000{,}000{,}007 = 20\).

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

{{PHOTO_ID:176939544513599}}

If the first road (between plants \(2\) and \(1\)) were blockaded, there would be a shipping capacity of \(5\) between plants \(1\) and \(4\), and a shipping capacity of \(1\) between plants \(1\) and \(3\) as well as between plants \(3\) and \(4\). Therefore, \(S_1 = 5 + 1 + 1 = 7\). Overall, \(S = [7, 20, 12]\), and the answer is \((7 * 20 * 12)\) modulo \(1{,}000{,}000{,}007 = 1{,}680\).

In the fourth town, \(S = [6, 4, 4, 6]\).

In the fifth town, \(S = [36, 41, 28, 21, 34, 29, 41]\).


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


Sample output:
Case #1: 0
Case #2: 20
Case #3: 1680
Case #4: 576
Case #5: 85240043


Solution: Let \(K\) be the sum of shipping capacities without any roads being blockaded. We can compute \(K\) in \(O(N \log(N))\) time by sorting and processing the roads in non-increasing order of capacity, while maintaining a [disjoint-set data structure](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) of groups of processing plants connected by roads processed so far (with each road initially in its own group). When road \(i\) is processed, it causes all plants in \(A_i\)'s group to become connected to all plants in \(B_i\)'s group, and must be a minimum-capacity road on the paths between those plants. Therefore, we can increment \(K\) by the product of \(C_i\), the number of plants in \(A_i\)'s group, and the number of plants in \(B_i\)'s group. The two groups should then be merged together.

We'll then want to compute, for each road \(i\), the amount \(Z_i\) by which the sum of shipping capacities would decrease if road \(i\) were blockaded (such that \(S_i = K - Z_i\)).

One possible approach involves two stages of dynamic programming, based on treating the network of plants and roads as a tree rooted at an arbitrary node (plant).

First, let \(D_{i,j}\) be the number of nodes within node \(i\)'s subtree which are reachable from node \(i\) using only edges (roads) with capacities greater than or equal to \(j\). We can compute \(D_{i,j}\) for all pairs \((i, j)\) by recursing through the tree — when a node \(i\) has an edge with capacity \(c\) to a child \(k\), we should increase \(D_{i,c'}\) by \(D_{k,c'}\) for each \(c'\) no greater than \(c\).

Next, let \(U_{i,j}\) similarly be the number of nodes *not* within node \(i\)'s subtree which are reachable from node \(i\) using only edges with capacities greater than or equal to \(j\). We'll recurse through the tree once again to compute these values — when a node \(i\) has an edge with capacity \(c\) to a child \(k\), we should set \(U_{k,c'}\) to equal \(U_{i,c'} + D_{i,c'} - D_{k,c'}\) for each \(c'\) no greater than \(c\).

Finally, let's consider each non-root node \(i\), such that the edge connecting node \(i\) and its parent corresponds to some road \(r\). The number of pairs of plants with shipping capacities greater than or equal to some capacity \(c\) whose routes include road \(r\) must be \(P_{i,c} = D_{i,c} * U_{i,c}\). Furthermore, the number of such pairs with shipping capacities equal to exactly \(c\) must be \(P_{i,c+1} - P_{i,c}\). We can therefore compute \(Z_i\) as the sum of \(c * (P_{i,c+1} - P_{i,c})\) over all capacities \(c\).

This algorithm takes a total of \(O(N (\log(N) + \max(C_{1..N})))\) time.

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


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

const int LIM = 800008, LIM2 = 20;
const int MOD = 1000000007;

class UnionFind {
  int N;
  int root[LIM], rank[LIM], size[LIM];

  void init(int _N) {
    N = _N;
    for (int i = 0; i < N; i++) {
      root[i] = i;
      size[i] = 1;
    }
  }

 public:
  UnionFind(int _N) {
    init(_N);
  }

  int find(int i) {
    if (root[i] != i) {
      root[i] = find(root[i]);
    }
    return root[i];
  }

  bool merge(int i, int j) {
    i = find(i);
    j = find(j);
    if (i == j) {
      return false;
    }
    if (rank[i] > rank[j]) {
      swap(i, j);
    }
    root[i] = j;
    size[j] += size[i];
    if (rank[i] == rank[j]) {
      rank[j]++;
    }
    return true;
  }

  int get_size(int i) {
    return size[find(i)];
  }
};

int N, ans, tot;
vector<pair<int, int>> adj[LIM];
int dynU[LIM][LIM2 + 2]; // dynU[i][j] = # of nodes reachable up from i (through edge to parent) via capacity >= j edges
int dynD[LIM][LIM2 + 2]; // dynD[i][j] = # of nodes reachable down from i via capacity >= j edges

// Populate dynD.
void recD(int i, int p) {
  for (int j = 1; j <= LIM2; j++) {
    dynD[i][j] = 1;
  }
  for (auto e : adj[i]) {
    int b = e.first, c = e.second;
    if (b == p) {
      continue;
    }
    recD(b, i);
    for (int j = 1; j <= c; j++) {
      dynD[i][j] += dynD[b][j];
    }
  }
}

// Populate dynU (and answer).
void recU(int i, int p) {
  for (auto e : adj[i]) {
    int b = e.first, c = e.second;
    if (b == p) {
      continue;
    }
    for (int j = 1; j <= LIM2; j++) {
      if (j <= c) {
        dynU[b][j] = (dynU[i][j] + dynD[i][j] - dynD[b][j] + MOD) % MOD;
      } else {
        dynU[b][j] = 0;
      }
    }
    int totE = tot, prev = 0;
    for (int j = c; j > 0; j--) {
      int curr = (long long)dynD[b][j]*dynU[b][j] % MOD;
      int cnt = (curr - prev + MOD) % MOD;
      int s = (long long)cnt*j % MOD;
      totE = (totE - s + MOD) % MOD;
      prev = curr;
    }
    ans = (long long)ans*totE % MOD;
    recU(b, i);
  }
}

int solve() {
  for (int i = 0; i < LIM; i++) {
    adj[i].clear();
  }
  // Input.
  cin >> N;
  vector<pair<int, pair<int, int>>> E;
  for (int i = 0; i < N - 1; i++) {
    int a, b, c;
    cin >> a >> b >> c;
    a--;
    b--;
    adj[a].push_back(make_pair(b, c));
    adj[b].push_back(make_pair(a, c));
    E.push_back(make_pair(c, make_pair(a, b)));
  }
  // Compute initial capacity sum.
  tot = 0;
  UnionFind U(N);
  sort(E.begin(), E.end());
  for (int i = (int)E.size() - 1; i >= 0; i--) {
    int a = E[i].second.first, b = E[i].second.second, c = E[i].first;
    tot = (tot + (long long)c*U.get_size(a)*U.get_size(b)) % MOD;
    U.merge(a, b);
  }
  // Compute reduced capacity sums via DP.
  ans = 1;
  recD(0, -1);
  recU(0, -1);
  return ans;
}

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