Official

F - AtCoder Express 3 Editorial by evima

We will consider a dynamic programming approach. However, it would be hard to do it based on a breadth-first search to find the shortest path from Station \(0\) to \(N\), since it would require us to maintain a complex state.

Thus, we will develop an easier way to find the shortest path, and then use it to build a dynamic programming solution.

For convenience, we will call stations and railways administered by Ko-soku red, and those administered by Jun-kyu blue.

First, we will develop a way to find the shortest path when the color of every station is fixed.

For now, let us assume that we never go back, that is, we only go from Station \(x\) to Station \(y\) where \(x < y\). In such a case, the following algorithm finds the shortest path from Station \(0\) to \(N\).

For \(i = 1, 2, \cdots, N-1\), in this order, update the value \((\alpha, \beta)\). Initially, \((\alpha, \beta) = (0, 0)\). (They are the lengths of the shortest paths to the last red station and the last blue station.)

  • If \(c_i\) is A and \(c_{i-1}\) is A: update \(\alpha\) to \(\alpha + 1\).
  • If \(c_i\) is A and \(c_{i-1}\) is B: update \(\alpha\) to \(\min(\alpha, \beta) + 1\).
  • If \(c_i\) is B and \(c_{i-1}\) is A: update \(\beta\) to \(\min(\beta, \alpha) + 1\).
  • If \(c_i\) is B and \(c_{i-1}\) is B: update \(\beta\) to \(\beta + 1\).
  • (If \(i = 1\), assume \(c_{i-1}\) to be A or B; whichever you choose, the value of \((\alpha, \beta)\) will be the same.

Then, the length of the shortest path to Station \(N\) is \(\min(\alpha, \beta) + 1\).

Unfortunately, there are cases the shortest path does not satisfy the assumption. The figure below shows one such case where the shortest path has a backward move:

However, when the shortest path goes backward, it is always by Local Train. Additionally, such a backward move never happens twice in a row in the shortest path. This is because going backward by Red or Blue Train, or by Local Train twice in a row, would result in stopping at the same station twice or more, which should not happen in the shortest path.

The figure below transforms the network to make it easier to see this fact:

It means we can slightly modify the algorithm above to compute the shortest path, as follows:

For \(i = 1, 2, \cdots, N-1\), in this order, update the value \((\alpha, \beta)\). Initially, \((\alpha, \beta) = (0, 0)\). (They are the lengths of the shortest paths to the last red station and the last blue station.)

  • If \(c_i\) is A and \(c_{i-1}\) is A: update \(\alpha\) to \(\alpha + 1\).
  • If \(c_i\) is A and \(c_{i-1}\) is B: update \(\alpha\) to \(\min(\alpha, \beta) + 1\), then use the new value of \(\alpha\) to update \(\beta\) to \(\min(\beta, \alpha + 1)\).
  • If \(c_i\) is B and \(c_{i-1}\) is A: update \(\beta\) to \(\min(\beta, \alpha) + 1\), then use the new value of \(\beta\) to update \(\alpha\) to \(\min(\alpha, \beta + 1)\).
  • If \(c_i\) is B and \(c_{i-1}\) is B: update \(\beta\) to \(\beta + 1\).
  • (If \(i = 1\), assume \(c_{i-1}\) to be A or B; whichever you choose, the value of \((\alpha, \beta)\) will be the same.

Then, the length of the shortest path to Station \(N\) is \(\min(\alpha, \beta) + 1\).

The figure below shows the behavior of this algorithm on some network, where a number represents the length of the shortest path to that point. In the case to the left, the bold edge is traversed from left to right; in the case at the center, the bold edge is irrelevant; in the case to the right, the bold edge is traversed from right to left. In all of those cases, each of the red and blue layers is one-way, enabling us to sequentially update \((\alpha, \beta)\).

Now, let us use this algorithm to build a dynamic programming solution. We need the following information in the DP:

  • \(pos\), the index of the last station we have seen.
  • \(\alpha\), the length of the shortest path to the last red station.
  • \(\beta\), the shortest length of the shortest path to the last blue station.
  • \(pos\), the color of the last station \(pos\). (Use \(0\) for red and \(1\) for blue.)

Let \(dp[pos][\alpha][\beta][pre]\) store the number of ways to set the colors of the stations until Station \(pos\) leading to the state \((pos, \alpha, \beta, pre)\), where \(pre\) is \(0\) if the last station \(pos\) is red and \(1\) if it is blue.

Since we are allowed to let \(pre = 0\) in the beginning, we can assume that the initial state is \((pos, \alpha, \beta, pre) = (0, 0, 0, 0)\). Similarly to the algorithm above, the transitions from \((pos, \alpha, \beta, pre)\) are as follows:

If \(pre = 0\): (the last station is red)

  • if \(c_{pos+1}\) is A or ?, the state can transition to \((pos+1, \alpha + 1, \beta, 0)\);
  • if \(c_{pos+1}\) is B or ?, the state can transition to \((pos+1, \min(\alpha, \min(\beta, \alpha) + 1 + 1), \min(\beta, \alpha) + 1, 1) = (pos+1, \min(\alpha, \beta + 2), \min(\beta, \alpha) + 1, 1)\).

If \(pre = 1\): (the last station is blue)

  • if \(c_{pos+1}\) is A or ?, the state can transition to \((pos + 1, \min(\alpha, \beta) + 1, \min(\beta, \min(\alpha, \beta) + 1 + 1), 0) = (pos + 1, \min(\alpha, \beta) + 1, \min(\beta, \alpha + 2), 0)\);
  • if \(c_{pos+1}\) is B or ?, the state can transition to \((pos + 1, \alpha, \beta + 1, 1)\).

We can fill the table in this way to find the answer as the sum of \(dp[N-1][\alpha][\beta][pre]\) over all triples \((\alpha, \beta, pre)\) such that \(\min(\alpha, \beta) \lt K\), in \(O(N^3)\) time, which is not fast enough.

The above DP has \(O(N^3)\) states, which must be reduced to make it run fast enough.

Now, see the following figure, where \(\alpha = 4\) stays the same but \(\beta\) gets greater and greater: \(\beta = 4, 5, 6, 7, 8, 9, \dots\). We needed to store values for many pairs \((\alpha, \beta)\) because of such a case.

However, if the next station is red, \(\beta\) will be updated to \(6\), no matter how large it is. That is, as shown in the figure below, we can assume \(\beta = 6\) when \(\beta\) becomes \(6\) or greater, without affecting the final length of the shortest path.

More generally, the following holds:

  • If \(pre = 0\), we can assume \(\alpha = \beta + 2\) if \(\alpha \geq \beta + 2\), making \(-2 \leq \beta - \alpha \leq +1\) always hold.
  • If \(pre = 1\), we can assume \(\beta = \alpha + 2\) when \(\beta \geq \alpha + 2\), making \(-1 \leq \beta - \alpha \leq +2\) always hold.

It limits the pairs \((\alpha, \beta)\) to consider, and now our DP table has approximately \(8N^2\) states \((pos, \alpha, \beta, pre)\). In the actual implementation, we will have to maintain values as \(dp[pos][\alpha][\beta - \alpha + 2][pre]\) or in some other way. Now the problem is solved in \(O(N^2)\) time.

Summary of the solution:

  1. Let \(dp[pos][\alpha][\beta][pre]\) be the number of ways to set the colors of the stations until Station \(pos\) so that the length of the shortest path to the last red and blue stations are \(\alpha\) and \(\beta\), respectively, and the color of the last station is \(pre\). Try to find these values sequentially.
  2. We may need to move backward, but we only have to do so by Local Train, and at most once in a row, so we can update \((\alpha, \beta)\) accordingly, leading to an \(O(N^3)\) solution.
  3. Actually, we can assume \(\beta = \alpha + 2\) when \(\beta - \alpha \geq 2\), and assume \(\alpha = \beta + 2\) when \(\alpha - \beta \geq 2\), which reduces the number of states to \(O(N)\), leading to an \(O(N^2)\) solution.

Sample Implementations:

(By multiplying polynomial matrices with FFT, we can modify our \(O(N^2)\) solution so that it runs in \(O(N \log^2 N)\) time. However, it would be practically much slower than the original because of the huge constant factor.)

posted:
last update: