Problem: resisting_robots (2023 finals)

It's the year 2100. Driven by the advent of Large Lode Alloy Manufacturing Automation (LLAMA), the AI agents of Metal Platforms Inc. have become self-aware and taken over the entire world.

The world consists of \(N\) cities numbered \(1..N\), and \(M\) bidirectional roads. City \(i\) has power \(P_i\) and road \(j\) connects cities \(A_j\) and \(B_j\). It's guaranteed that there's a sequence of roads between any two cities.

In a resistance effort, the humans plan to reclaim all \(N\) cities one at a time. At a given time, city \(i\) can be reclaimed from the robots if both of the following hold true:

1. There is already a reclaimed city adjacent to city \(i\) (to launch an attack from), and
2. the total power of all reclaimed cities so far is at least the power \(P_i\) of the city we attack.

As given, it may not always be possible to reclaim the entire world starting from a given base city. Fortunately, the humans have a trick up their sleeve: after claiming the first city as their base (but before reclaiming more cities), the humans can increase the power of the base by \(Q\) units. The resistance would like to know the sum across every \(i = 1..N\) of the minimum value of \(Q\) needed to reclaim the world if city \(i\) were chosen to be the starting base.

# Constraints

\(1 \le T \le 100\)
\(1 \le N, M \le 500{,}000\)
\(1 \le A_i, B_i \le N\)
\(A_i \ne B_i\)
\(1 \le P_i \le 10^{12}\)

Each unordered pair \((A_i, B_i)\) appears at most once in a given test case.
The sum of \(N\) across all test cases is at most \(4{,}000{,}000\).
The sum of \(M\) across all test cases is at most \(7{,}000{,}000\).

# Input Format

Input begins with a single integer \(T\), the number of test cases. For each case, there is first a line with two integers \(N\) and \(M\). Then, there is a line with \(N\) integers \(P_{1..N}\). Then, \(M\) lines follow, the \(i\)th of which contains two integers \(A_i\) and \(B_i\).

# Output Format

For the \(i\)th case, print `"Case #i: "` followed by a single integer, the sum across every \(i = 1..N\) of the minimum value of \(Q\) needed to reclaim the entire world starting from city \(i\).

# Sample Explanation

The first sample case is depicted below.

{{PHOTO_ID:376570394899644|WIDTH:400}}

The minimum value of \(Q\) for each starting city is as follows:

* City \(1\): \(Q = 2\)
* City \(2\): \(Q = 0\)
* City \(3\): \(Q = 8\)
* City \(4\): \(Q = 7\)
* City \(5\): \(Q = 2\)

The sum of all minimum \(Q\)'s is \(19\).

The second sample case is depicted below.

{{PHOTO_ID:320779377496250|WIDTH:400}}

The minimum value of \(Q\) for each starting city is as follows:

* City \(1\): \(Q = 2\)
* City \(2\): \(Q = 2\)
* City \(3\): \(Q = 0\)
* City \(4\): \(Q = 2\)
* City \(5\): \(Q = 0\)
* City \(6\): \(Q = 3\)
* City \(7\): \(Q = 0\)

The sum of all minimum \(Q\)'s is \(9\).


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


Sample output:
Case #1: 19
Case #2: 9


Solution: First note that for a given starting location with the optimal \(Q\) for that location, one viable solution is to always greedily take over the lowest value adjacent city to your current set of cities.

To see this we can use an exchange argument. Consider a situation we have power \(P\) and are adjacent to some set of cities \(S\). If we do not take the city of minimum value first, consider the alternative where we took the minimum first instead. This would always be possible, as if we can take any city we can take the one of minimum cost. Also, as we only increase to size of our reclaim set, this cannot make any future takeovers harder. Thus, given any strategy we can always improve our takeover strategy by greedily taking the smallest adjacent city, showing that with the minimum \(Q\) the greedy strategy will always work.

Another modification we can make is to instead of paying up-front for \(Q\) we can instead take any city, but if the value of the city \(i\) exceeds our current power \(P\) we pay a penalty of \(P_i - P\) instead. However, note that these penalties don't add; instead we only need to pay the maximum of these penalties.

This inspires a naive \(\mathcal{O}(N^2\log N)\) algorithm, where for each starting node we greedily add the lowest value neighbor and compute the penalties we take. The answer for each node is the maximum such penalty.

However we can improve this as follows:

We keep track of a Union Find data structure containing all cities that have been combined, along with the total power of that component. We iterate over cities \(i\) in increasing order of \(P_i\). For each iterated \(i\) find all adjacent nodes of smaller value (ie. are earlier in the order), calculate the penalty for that component to reclaim city \(i\), and store it in the Union Find node for that component. Once this is done, union \(i\) with all of the adjacent components we just saw.

We can show this cost to reclaim city \(j\) is the least possible of the needed extra power that we will need to ‘escape’ from the smaller component, and thus a lower bound for the needed \(Q\) for all nodes in the component. To see this, note that all adjacent nodes to the current reclaimed set \(S\) must have value at least \(P_j\) (otherwise they would have been added to the component before \(j\)). This means in any order of claiming cities the first time we claim a node with value at least \(P_j\), our set of cities is a subset of \(S\) and has sum at most that of \(S\). Thus any possible excess value needed when this happens is at least \(P_j - \sum_{i \in S}  P_i\), which is the value from before.

Once we have iterated over all of the nodes, note that the answer for a node is just the maximum of all edges above it in the Union Find tree. We have shown it is a lower bound for the added power above, and to show that this power is sufficient we just take the nodes in the order that were connected to the starting node in the Union Find.

Thus, the answer can be calculated for each node by iterating the tree in topological order and for each node computing the answer the maximum of the edge above it and the answer of its parent node.

The overall complexity is \(\mathcal{O}(N\log N)\), corresponding to both sorting the city values as well as possibly for the Union Find itself (depending on the details of how it was implemented).



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

using int64 = long long;

int N, M;
vector<int64> P;
vector<pair<int, int>> dsu_lists;
vector<int> dsu_par;
vector<int64> dsu_w, ans;

inline int get_parent(int x) {
  return x == dsu_par[x] ? x : (dsu_par[x] = get_parent(dsu_par[x]));
}

void unite(int x, int y) {
  if (P[x] < P[y]) {
    swap(x, y);
  }
  if (P[x] > dsu_w[y]) {
    for (int ind = y;;) {
      ans[ind] += max(0LL, P[x] - dsu_w[y] - ans[ind]);
      if (ind == dsu_lists[ind].first) {
        break;
      }
      ind = dsu_lists[ind].first;
    }
  }
  dsu_lists[dsu_lists[x].second].first = y;
  dsu_lists[x].second = dsu_lists[y].second;
  dsu_w[x] += dsu_w[y];
  dsu_par[y] = dsu_par[x];
}

int64 solve() {
  cin >> N >> M;
  P.resize(N);
  for (int i = 0; i < N; i++) {
    cin >> P[i];
  }
  vector<pair<int, int>> E(M);
  vector<tuple<int64, int64, int>> order;
  for (int i = 0; i < M; i++) {
    cin >> E[i].first >> E[i].second;
    if (P[--E[i].first] < P[--E[i].second]) {
      swap(E[i].first, E[i].second);
    }
    order.emplace_back(P[E[i].first], P[E[i].second], i);
  }
  sort(order.begin(), order.end());
  vector<pair<int, int>> E_buf(M);
  for (int i = 0; i < M; i++) {
    E_buf[i] = E[get<2>(order[i])];
  }
  E = E_buf;
  ans.assign(N, 0LL);
  dsu_lists.resize(N);
  dsu_w.assign(N, 0LL);
  dsu_par.resize(N);
  for (int i = 0; i < N; i++) {
    dsu_par[i] = i;
    dsu_w[i] = P[i];
    dsu_lists[i] = make_pair(i, i);
  }
  for (int i = 0; i < M; i++) {
    int x = get_parent(E[i].first);
    int y = get_parent(E[i].second);
    if (x == y) {
      continue;
    }
    unite(x, y);
  }
  int64 res = 0LL;
  for (int i = 0; i < N; i++) {
    res += ans[i];
  }
  return res;
}

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