OZ - 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:

TBD

PREREQUISITES:

None

PROBLEM:

You’re given a binary string S.
You can swap adjacent characters.
A prefix of S is called good if it contains at least as many ones as zeros.

Find the maximum possible number of good prefixes possible, and also the minimum number of swaps needed to reach this maximum.

EXPLANATION:

Finding the maximum possible number of good prefixes is fairly simple: the natural greedy algorithm of simply placing all ones at the start and zeros afterwards (equivalently, sorting in descending order) results in the maximum possible number of good prefixes.

To see why: if S_i = 0 and S_{i+1} = 1, it can be proved that swapping them doesn’t make the answer worse (in particular, the only prefix that gets affected at all is the one of length i, and it now has one extra 1 compared to a 0 so it’s clearly doing not worse.)

In particular, if the string has C ones, the maximum number of prefixes is exactly \min(2C, N).

We now need to find the minimum number of swaps needed to reach the maximum number of good prefixes.

Note that it’s possible to reach the maximum count without having all ones as a prefix: for example S = \texttt{101} has all its prefixes be good, so no swaps are needed.

However, it can instead be proved that in an optimal configuration, if the prefix of length i is good then the prefixes of all lengths \le i must be good.

Proof

The idea is similar to how we proved that that starting with all ones is optimal to maximize the count.

Suppose the prefix of length i is not good while the one of length j is, where i \lt j.
If there are multiple such i, choose the leftmost one among them, and if there are multiple such j for this i, again choose the leftmost one among them.

Then, it can be seen that we must have S_i = 0 and S_{j} = 1.
(If S_i = 1, then the prefix of length i-1 wouldn’t be good either, contradicting minimality of i.)

In particular, there surely exists an occurrence of 1 in the range [i+1, j].
Let k be the leftmost occurrence of 1 in this range.
Repeatedly swap this one left till it reaches position i, and observe that now, prefix i will become good while the status of all other prefixes doesn’t change, so we obtain a strictly better result.

In particular, once we know the maximum number of good prefixes, we also know exactly which ones to make good.

The above proof also gives us a fairly simple greedy algorithm to solve the problem - while there exists an index i such that the length i prefix is not good but the length i+1 prefix is good, swap i with i+1.

However, directly simulating this will be too slow, since in the worst case we might need \mathcal{\Theta}(N^2) swaps to reach the final state.
It is, however, not too hard to speed this algorithm up, as follows.

Iterate i from 1 to N.
When processing index i,

  • If the prefix of length i is already good, nothing needs to be done.
  • Otherwise, look at indices \gt i.
    • If all these positions contain 0's, nothing further can be done.
    • If at least one occurrence of 1 is present however, we can make the prefix of length i good.
      This will be done by choosing the leftmost occurrence of 1 at an index \gt i, and repeatedly performing swaps to bring it to position i.
      Observe that this is equivalent to what the slower greedy algorithm we had above will do, just compressing several steps into one.

This can already be implemented to run in \mathcal{O}(N\log N) with some data structure use, but there exists an even simpler (to implement) solution.

Observe that what the above algorithm is doing, is simply ensuring that for every k, the k-th occurrence of 1 in S is placed somewhere at indices 1, 2, \ldots, 2k+1.
This is because if this is not done, then the prefix of length 2k+1 cannot be good.

Thus, we can easily determine the final state of every occurrence of 1 in S as follows:

  • If the k-th 1 appears at position i_k then,
    • If i_k \le 2k+1, this 1 won’t move at all.
      This has a cost of 0.
    • Otherwise, this 1 will move to exactly position 2k+1.
      This has a cost of i_k - (2k+1).

This allows for a simple linear-time implementation.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for _ in range(int(input())):
    n = int(input())
    s = input()
    
    ones = s.count('1')
    mx_count = min(2*ones, n)
    
    r = 0
    ans = 0
    for i in range(n):
        if s[i] == '1':
            # if >= r, bring to r
            if i >= r: ans += i - r
            r += 2
    
    print(mx_count, ans)