REDBLUESW - Editorial

editorial
, , , ,
1

PROBLEM LINK:

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

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

Dynamic programming, prefix sums

PROBLEM:

You’re given a permutation of the integers [1, N]. Each element also has a color - either red or blue.
You can perform the following operation:

  • Choose i such that elements i and i+1 have different colors; and the red element is larger than the blue element.
    Then swap these two elements.

Count the number of reachable permutations.

EXPLANATION:

Each adjacent swap involves one red element and one blue element.
Thus, we are completely unable to change the relative order of red elements within themselves, nor the order of blue elements within themselves.

One way to think of this is as follows.
Suppose there are R red and B blue elements (where R+B=N.)

We can first place all R red elements in a line.
Then, for each blue element, there are R+1 positions where it can be placed: either between two red elements, or before the first one, or after the last one.
Observe that if we decide for each blue element, which of the R+1 positions it is going to take - the the final permutation is also fixed, since the relative positions of the blue elements are themselves fixed.

This fact allows us to use dynamic programming to count the number of reachable permutations.

Let’s number the buckets to place blue elements as 0, 1, \ldots, R from left to right, where:

  • 0 is the bucket before the first element.
  • For each 1 \le i \le R, bucket i is immediately after the i-th red element.

Let’s also number the blue elements 1, 2, \ldots, B from left to right.

Each blue element is initially in one of the buckets 0, 1, \ldots, R.
Let s_i be the initial bucket of the i-th element.

If we want to move this element to a different bucket, then it must be able to swap across all red elements in its path - meaning it must be smaller than all of them.
This means that for each blue element, there’s a certain range of buckets it can reach.

Let l_i and r_i denote the leftmost/rightmost buckets that the i-th blue element can reach.
These can be found in \mathcal{O}(N) for each i by simply iterating over the red elements; which is fine since the constraints allow \mathcal{O}(N^2) anyway.

Define dp(i, j) to be the number of ways of placing the first i blue elements, such that the i-th blue element is placed in the j-th space.
Of course, if j is not in the range [l_i, r_i] we can immediately say dp(i, j) = 0; so we only care about the case where l_i \le j \le r_i.

Now, if the i-th blue element is placed in the j-th bucket, then the (i-1)-th blue element cannot appear after it - and hence must be present in at most the j-th bucket as well.
Summing up over all possible cases, we see that

dp(i, j) = dp(i-1, 0) + dp(i-1, 1) + \ldots + dp(i-1, j)

Of course, this is under the assumption that l_i \le j \le r_i, with the dp value being 0 otherwise.

The final answer is dp(B, 0) + dp(B, 1) + \ldots + dp(B, R), i.e. we sum up across all possible ways of placing the last element.

There are \mathcal{O}(N^2) states for this DP and \mathcal{O}(N) transitions from each, leading to a complexity of \mathcal{O}(N^3).
To speed it up, observe that dp(i, j) equals either 0 or the j-th prefix sum of dp(i-1, \cdot).
So, by simply maintaining these prefix sums, we are able to perform transitions in \mathcal{O}(1) time, leading to a complexity of \mathcal{O}(N^2) overall which is fast enough for the given limits.

TIME COMPLEXITY:

\mathcal{O}(N^2) per testcase.

CODE:

Editorialist's code (PyPy3) 
mod = 998244353
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))

    red = [0]
    for i in range(n):
        if b[i] == 0:
            red.append(a[i])
    red.append(0)
    
    lt, rt = [], []
    cur = 0
    for i in range(n):
        if b[i] == 1:
            lt.append(cur)
            rt.append(cur)

            while a[i] < red[lt[-1]]:
                lt[-1] -= 1
            while a[i] < red[rt[-1]+1]:
                rt[-1] += 1
        else:
            cur += 1
    
    sz = len(red)
    dp, ndp = [0]*sz, [0]*sz
    dp[0] = 1
    for i in range(len(lt)):
        l, r = lt[i], rt[i]

        pref = 0
        for j in range(sz):
            pref = (pref + dp[j]) % mod

            if l <= j <= r:
                ndp[j] = pref
            else:
                ndp[j] = 0
        dp, ndp = ndp, dp
    print(sum(dp) % mod)
1 Like
2

nice problem