DIFFVAL - Editorial

PROBLEM LINK:

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

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

None

PROBLEM:

You’re given a binary string S and an integer K.
Does there exist a rearrangement of S such that S_i \ne S_{i+K} for all 1 \le i \le N-K?

EXPLANATION:

Suppose we fix what S_1 is; say we fix S_1 = 0.
Observe that this then fixes several other characters, namely:

  • We want S_1 \ne S_{K+1}, so S_{K+1} must be 1.
  • We want S_{K+1} \ne S_{2K+1}, so S_{2K+1} must be 0.
  • We want S_{2K+1} \ne S_{3K+1}, so S_{3K+1} must be 1.
    \vdots

In general, after fixing S_1, the characters at all positions of the form S_{xK+1} will also be uniquely fixed.
Further, observe that the sequence S_1, S_{K+1}, S_{2K+1}, \ldots will alternate between zeros and ones.
Specifically,

  • If this sequence has an even length, then it will use an equal number of zeros and ones no matter whether we set S_1 = 0 or S_1 = 1.
  • If this sequence has an odd length, then it will use either one extra 0, or one extra 1, depending on what S_1 is.

Observe that the above analysis works for any starting point, not just index 1.
That is, if we choose a starting position p, then fixing the value of S_p will automatically decide the values of all indices of the form (xK + p); and the values at these indices must alternate.
So, depending on whether the sequence length is even or odd, we’ll need to use either an equal number of zeros and ones; or one extra instance of one of them.


There are K possible starting positions: one for each of the indices 1, 2, 3, \ldots, K.
Further indices aren’t starting positions, since their values will be determined once the first K characters of S are determined.

So, we have K possible “chains”: one from each starting position.
Recall that for each chain,

  • If its length is even, then it uses an equal number of zeros and ones.
  • If its length is odd, then it uses either one extra 0 or one extra 1.

Observe that for even chains it doesn’t really matter what we do since we use an equal number of zeros/ones anyway; so we only really need to focus on the odd chains.

Suppose we have X odd chains to deal with.
Also let c_0 and c_1 denote the number of zeros and ones available to us.

For each odd chain, it must contain either an extra zero or an extra one.
Suppose r of the odd chains contain an extra zero; so that (X-r) of them contain an extra one.
For this to be possible:

  • r \le c_0 must be true, obviously - we do need at least r zeros.
  • Similarly, X-r \le c_1 must be true.
  • After deciding the extra zeros/ones, each chain now has effectively an equal number of zeros and ones.
    This means the remaining number of zeros and ones must be equal, i.e.
    c_0 - r = c_1 - (X-r)

If all these conditions are satisfied for some r, a valid rearrangement will always exist.

Why?

Let the odd-length chains be the ones starting at indices i_1, i_2, \ldots, i_X.
Suppose we want r of them to have an extra 0.
We can then simply place zeros at positions i_1, i_2, \ldots, i_r, and then place ones at positions i_{r+1}, \ldots, i_X.

Note that because we ensured (c_0 - r) = c_1 - (X-r), we’re now left with an equal number of zeros and ones, which need to be placed in the remaining positions.

This allows for the following:

  • For each even-length chain, just make it 010101\ldots 01.
  • For each odd-length chain, make it 0101\ldots 010 or 1010\ldots 101 depending on whether it starts with a 0 or a 1 (which we decided above).

Since we have an equal number of zeros and ones remaining, this will exactly fill up all remaining positions of the string while ensuring S_i \ne S_{i+K}, completing our construction.

So, after computing c_0, c_1, X, we only need to check if a valid r exists.
This can be done by just iterating through all choices of r from 0 to X, or even in \mathcal{O}(1) with a bit of math (though the overall solution remains linear anyway.)

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n, k = map(int, input().split())
    s = input()
    
    c0, c1 = s.count('0'), s.count('1')
    odd = 0
    for i in range(k):
        sz = 0
        for j in range(i, n, k): sz += 1
        odd += sz%2
    
    ans = 'No'
    for r in range(odd+1):
        if r <= c0 and odd-r <= c1 and c0-r == c1-(odd-r):
            ans = 'Yes'
    print(ans)