Shortest Path with Color Transformations: Solving the JOI 2021 Final Robot Problem
We are given a connected undirected graph (or possibly disconnected) with N vertices and M edges. Each edge e connects u and v, has an integer color c, and a cost p for traversing it. However, at any vertex, if multiple incident edges share the same color, the robot cannot decide which path to take. We can pay to modify colors: either change the color of a single edge to a unique one (pay p), or change all other incident edges of a given color at the current vertex to a new color (pay the sum of their traversal costs minus the cost of the edge we intend to keep). The goal is to go from node 1 to node N with minimum total modification cost, assuming we can traverse an edge for free once its color is unique or after resolving conflicts.
First, note that if there is exactly one incident edge with the needed color at the current vertex, traversal costs 0. More generally, from a vertex u, if we want to traverse edge (u,v) with color c and cost p, we have two options to make the path unambiguous:
- Pay p and recolor this edge to a color not appearing at u.
- Pay sum_c - p, where sum_c is the total cost of all incident edges at u that have color c. This recolors all other edges of color c, leaving only the desired edge (u,v) with color c. Thus the effective cost to go from u to v along this edge is min(p, sum_c - p). A naive approach would build a graph with these direct edge weights and run Dijkstra. How ever, this fails because altering edges on the current side has consequences on the opposite vertex v. When we pay sum_c - p to clear other edges at u, we effectively make the color c available at v as well, possibly eliminating the need to pay again for modifications at v when continuing with color c. This "after‑effect" must be captured.
To handle it correctly, we introduce virtual (shadow) nodes that represent the state after a color has been cleaned on one side. For every vertex x and every color col that appears among its incident edges, create a virtual node VID(x, col). Let sum(x, col) be the total cost of edges incident to x having color col.
Now for each original edge (u, v, col, p) we insert the following directed edges into our model (and symmetrically for the opposite direction):
- u → v with weight min(p, sum(u, col) - p)
- u → VID(v, col) with weight 0
- VID(u, col) → v with weight sum(u, col) - p
Interpretation: The direct edge u→v corresponds to immediately paying for the modification at u (either recoloring this single edge or clearing the other edges) and going to v. The path u → VID(v, col) → … models the scanario where we clear the other edges of color col at u (cost already accounted as sum(u,col)-p, which will be paid when we later use the VID(u,col)→v edge or when we traverse an edge from VID(u,col) to somewhere else). Actually, the three edges together ensure that the optimal cost is captured: if we want to reuse the cleared color col at v, we can go u → VID(v, col) for 0, and then from VID(v, col) we can go to other vertices via edges that expect color col to be available at v (the edge from VID(v, col) to some neighbor will pay the necessary clearing cost on the v side). The edge VID(u, col) → v represents that after we have already paid to clear other color‑col edges at u, we can enter v with that color ready. Combined with the costs, any valid sequence of modifications corresponds to a path in this directed graph, and the shortest path yields the optimal answer.
After building the graph with at most N + 2M vertices and at most 6M directed edges, we run Dijkstra from node 1. If node N is unreachable, output -1.
The overal time complexity is O((N+M) log (N+M)).
Below is an implementation in C++.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = 0x3f3f3f3f3f3f3f3fLL;
struct EdgeInfo {
int to, color;
ll cost;
};
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
int n, m;
cin >> n >> m;
vector<vector<EdgeInfo>> adj(n + 1);
for (int i = 0; i < m; i++) {
int u, v, c; ll p;
cin >> u >> v >> c >> p;
adj[u].push_back({v, c, p});
adj[v].push_back({u, c, p});
}
int nodes = n;
// assign IDs for virtual nodes VID(vertex, color)
vector<unordered_map<int, int>> vid(n + 1);
for (int u = 1; u <= n; u++) {
for (auto &e : adj[u]) {
int col = e.color;
if (!vid[u].count(col)) {
vid[u][col] = ++nodes;
}
}
}
// total cost per color per vertex
vector<unordered_map<int, ll>> sum_cost(n + 1);
for (int u = 1; u <= n; u++) {
for (auto &e : adj[u]) {
sum_cost[u][e.color] += e.cost;
}
}
// Build directed graph
vector<vector<pair<int, ll>>> g(nodes + 1);
for (int u = 1; u <= n; u++) {
for (auto &e : adj[u]) {
int v = e.to, col = e.color;
ll p = e.cost;
ll sum_c = sum_cost[u][col];
// Option 1: direct modification
ll w1 = min(p, sum_c - p);
g[u].push_back({v, w1});
// Option 2: go to the virtual node of v for this color
g[u].push_back({vid[v][col], 0});
// Option 3: from virtual node of u, enter v
g[vid[u][col]].push_back({v, sum_c - p});
}
}
// Dijkstra
vector<ll> dist(nodes + 1, INF);
priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<>> pq;
dist[1] = 0;
pq.push({0, 1});
while (!pq.empty()) {
auto [d, u] = pq.top(); pq.pop();
if (d != dist[u]) continue;
for (auto &[v, w] : g[u]) {
if (dist[v] > d + w) {
dist[v] = d + w;
pq.push({dist[v], v});
}
}
}
if (dist[n] == INF) cout << -1 << '\n';
else cout << dist[n] << '\n';
return 0;
}