OTREO - Editorial

editorial
, ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Medium

PREREQUISITES:

Dynamic Programming

PROBLEM:

For an array A containing only ones and twos, define f(A) as follows:

  • Let P be a permutation of [1, N-1].
  • For each i, if A_{P_i} = 1 and A_{P_i + 1} = 2, swap them; otherwise do nothing.
  • f(A) is the count of distinct final arrays under this process.

Given an array A, compute f(A[i\ldots N]) for each of its suffixes.

EXPLANATION:

Let’s analyze which final arrays are reachable.

Note that each swap essentially allows us to “move” an occurrence of 2 one step to the left.
The final array is also uniquely determined by the positions of the 2's in it, so let’s focus on that.

Let the initial positions of the 2's be i_1, i_2, \ldots, i_K from left to right.
What can their final positions be?

Let’s look at the occurrences in order.
Consider the first 2, at index i_1.
Then,

  • If i_1 = 1, this 2 is already at the leftmost position; and cannot move further.
    So it has only one option.
  • Otherwise, there are i_1 - 1 positions before it.
    It can end up in any of those positions – and also must end up in one of those positions (since it must move at least once, when i-1 appears in the permutation).
    So, it has i_1 - 1 options here.

Next, let’s look at i_2.
Just as above, this 2 must also move left a few steps: it can definitely reach every position in [i_1+1, i_2-1].

Further, note that this 2 can never end up at a position that’s \leq i_1 - 1.
This is because, for this to happen, we require the swap (i_1-1, i_1) to happen; but for that to even be possible, the 2 that was originally at i_1 must have moved to a position \lt i_1-1 itself, and doing so needs the (i_1-1, i_1) swap.
Since P is a permutation, each swap can be made at most once, so it’s impossible for the 2 at i_2 to end up at a position that’s \leq i_1-1.

However, there are still a couple of other options:

  • First, it’s possible to end up at index i_1 itself, though that’s only possible if the 2 at this position moved away first.
  • Second, only when i_1+1 = i_2, it’s also possible for this 2 to not move at all; by attempting to perform the swap (i_2-1, i_2) before moving the previous 2 away.

So, things are already getting a bit complicated: we have i_2 - i_1 - 1 options, but one extra (ending at i_1) that’s only available when the 2 at i_1 has moved away; and we also have one extra option specifically when i_1+1 = i_2.

For every index beyond i_2, the analysis is exactly the same as it was with i_2: there are several options for where it can end up, but one of them is only available if the previous 2 moved away, and one more (not moving) is only available if the previous 2 is initially adjacent to it.

This relationship suggests a dynamic programming approach; and it’s not too hard to see what the states are: dp[i][j] can denote the answer after processing the first characters, with j = 0/1 denoting whether the i-th two moved or not; 0 meaning it moved.

As for transitions:

  • If A_i = 1 we have dp[i] = dp[i-1].
  • If A_i = 2 then we have a couple of cases.
    • If this is the first occurrence of 2, then dp[i] will either equal [0, 1] or [i-1, 0], depending on whether i = 1 or not (so if it’s forced to move or not).
    • Otherwise, let k \lt i be the previous occurrence of 2.
    • If k = i-1, dp[i][0] = dp[i-1][0] and dp[i][1] = dp[i-1][0] + dp[i-1][1].
      This is because i can move only if i-1 moves; and if i doesn’t move then we can do anything till i-1.
    • If k \lt i-1, then dp[i][0] = dp[k][0] \cdot (i-k) + dp[k][1] \cdot (i-k-1)
      dp[i][1] will be 0, since this 2 will always have to move.

This DP runs in \mathcal{O}(N) time.
Of course, the issue here is that this takes \mathcal{O}(N) time for the entire array (so we can’t afford to run it for each suffix), and it computes the answer for each prefix, while we want the answer for each suffix.
The next step is hence to think about how to adapt this solution.

A bit of thought shows that modifying this solution to meet our needs is in fact quite easy: if we look at the ones instead of the twos, they move in exactly the opposite way.
So, all that needs to be done is to reverse the string and flip every 1 to 2 and vice versa; then run the above DP.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for i in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    a = a[::-1]

    mod = 998244353
    prv = -1
    dp = [0, 1]
    ans = []
    for i in range(n):
        if a[i] == 1:
            if prv == -1:
                if i == 0: dp = [0, 1]
                else: dp = [i, 0]
            elif prv == i-1:
                s = (dp[0] + dp[1]) % mod
                dp = [dp[0], s]
            else:
                d = i-prv-1
                s = (dp[0]*(d+1) + dp[1]*d) % mod
                dp = [s, 0]
            prv = i
        ans.append((dp[0] + dp[1]) % mod)
    ans = ans[::-1]
    print(*ans)
1 Like