JOI 2021 - Robot

Explanation

Lemma: We can always recolour an edge $(u, v)$ so that its colour is different from those of the other edges incident to $u$ and $v$.

Proof: If we take away $(u, v)$, then there are at most $M - 1$ edges incident to $u$ and $v$. Since there are $M$ colours in total, there must exist some colour such that none of those edges have it (by the pigeonhole principle). Therefore, we can simply recolour $(u, v)$ to that colour.

This means we can define the cost to traverse an edge as:

It may be tempting to use Dijkstra's algorithm at this point, but that will unfortunately result in a wrong answer. Consider the following:

• The robot goes on the path $u \rightarrow v \rightarrow w$.
• $(u, v)$ and $(v, w)$ have the same colour.
• It's optimal to recolour $(u, v)$ when going $u \rightarrow v$, but it's optimal to recolour the same-coloured neighbours of $(v, w)$ when going $v \rightarrow w$.

In this case, using Dijkstra's algorithm naively will count the cost to recolour $(u, v)$ twice in our answer!

To amend this, let's define two DP arrays instead of just one:

• $\texttt{dp1}[i]$ is the minimum cost to get to node $i$.
• $\texttt{dp2}[i][c]$ is the minimum cost to get to node $i$ if:

  • The robot just traversed a $c$-coloured edge to get to node $i$.
  • We're meant to have recoloured that edge, but haven't yet.
  • The robot will definitely leave node $i$ via another $c$-coloured edge and we will recolour all of its same-coloured neighbours.

We can then use Dijkstra's algorithm and some casework to solve this problem in amortized $\mathcal O((N + M) \log N)$ time.

Implementation

Time Complexity: $\mathcal O((N + M) \log N)$

Copy
#include <bits/stdc++.h>
typedef long long ll;
using namespace std;

const ll INF = 1e18;

struct Edge {
	int to, c;
	ll p;
};