Trees and Segments (Codeforces 1858D)
Tags: Greedy, two pointers

The teachers of the Summer Informatics School decided to plant $$$n$$$ trees in a row, and it was decided to plant only oaks and firs. To do this, they made a plan, which can be represented as a binary string $$$s$$$ of length $$$n$$$. If $$$s_i = 0$$$, then the $$$i$$$-th tree in the row should be an oak, and if $$$s_i = 1$$$, then the $$$i$$$-th tree in the row should be a fir.The day of tree planting is tomorrow, and the day after tomorrow an inspector will come to the School. The inspector loves nature very much, and he will evaluate the beauty of the row as follows:  First, he will calculate $$$l_0$$$ as the maximum number of consecutive oaks in the row (the maximum substring consisting of zeros in the plan $$$s$$$). If there are no oaks in the row, then $$$l_0 = 0$$$.  Then, he will calculate $$$l_1$$$ as the maximum number of consecutive firs in the row (the maximum substring consisting of ones in the plan $$$s$$$). If there are no firs in the row, then $$$l_1 = 0$$$.  Finally, he will calculate the beauty of the row as $$$a \cdot l_0 + l_1$$$ for some $$$a$$$ — the inspector's favourite number. The teachers know the value of the parameter $$$a$$$, but for security reasons they cannot tell it to you. They only told you that $$$a$$$ is an integer from $$$1$$$ to $$$n$$$.Since the trees have not yet been planted, the teachers decided to change the type of no more than $$$k$$$ trees to the opposite (i.e., change $$$s_i$$$ from $$$0$$$ to $$$1$$$ or from $$$1$$$ to $$$0$$$ in the plan) in order to maximize the beauty of the row of trees according to the inspector.For each integer $$$j$$$ from $$$1$$$ to $$$n$$$ answer the following question independently:   What is the maximum beauty of the row of trees that the teachers can achieve by changing the type of no more than $$$k$$$ trees if the inspector's favourite number $$$a$$$ is equal to $$$j$$$? 

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 1000$$$) — the number of test cases.The first line of each test case contains two integers $$$n$$$ and $$$k$$$ ($$$1 \le n \le 3000$$$, $$$0 \le k \le n$$$) — the number of trees in the row and the maximum number of changes.The second line contains a string $$$s$$$ of length $$$n$$$, consisting of zeros and ones — the plan description.It is guaranteed that the sum of all $$$n$$$ values for all test cases does not exceed $$$3000$$$.

Output: For each test case, print $$$n$$$ integers, the $$$j$$$-th ($$$1 \le j \le n$$$) of which is the maximum beauty of the row of trees after no more than $$$k$$$ changes if $$$a = j$$$ is used to calculate the beauty.

Note: In the first test case no changes are allowed, so $$$l_0 = 0$$$ and $$$l_1 = 3$$$ always hold. Thus, regardless of the value of $$$a$$$, the beauty of the row of trees will be $$$3$$$.In the second test case for $$$a \in \{1, 2\}$$$ the teachers can, for example, change the plan $$$s$$$ to $$$0111$$$ (by changing $$$s_4$$$), and for $$$a \in \{3, 4\}$$$ — to $$$0010$$$ (by changing $$$s_2$$$). In this case, the beauty of the row for each $$$a$$$ is calculated as follows:  For $$$a = 1$$$: $$$l_0 = 1$$$, $$$l_1 = 3$$$. The beauty of the row is $$$1\cdot 1 + 3 = 4$$$.  For $$$a = 2$$$: $$$l_0 = 1$$$, $$$l_1 = 3$$$. The beauty of the row is $$$2\cdot 1 + 3 = 5$$$.  For $$$a = 3$$$: $$$l_0 = 2$$$, $$$l_1 = 1$$$. The beauty of the row is $$$3\cdot 2 + 1 = 7$$$.  For $$$a = 4$$$: $$$l_0 = 2$$$, $$$l_1 = 1$$$. The beauty of the row is $$$4\cdot 2 + 1 = 9$$$. It can be shown that the changes described above are optimal for all $$$a$$$ from $$$1$$$ to $$$4$$$.

Sample Inputs:
5
3 0
111
4 1
0110
5 0
10000
6 2
101101
7 1
0001101

Sample Outputs:
3 3 3 
4 5 7 9 
5 9 13 17 21 
6 9 13 17 21 25 
7 10 13 17 21 25 29

Solution (Python):
def search_on_pref(j, zam, dl, x, sr, k):
    while j >= 0:
        dl += 1
        if s[j] != sr:
            zam += 1
        if zam <= k:
            x[zam] = max(x[zam], dl)
        j -= 1
    return x

def main(n, k, s):
    ans = [-1] * (n + 1)
    a = [0] * (k + 1)
    b = [0] * (k + 1)
    for i in range(n):
        a = search_on_pref(i, 0, 0, a, '0', k)
        b = search_on_pref(i, 0, 0, b, '1', k)
        for ind in range(1, k + 1):
            a[ind] = max(a[ind], a[ind - 1])
            b[ind] = max(b[ind], b[ind - 1])
        l = i
        zam = 0
        dl = 0
        while l < n:
            if zam <= k:
                ans[a[k - zam]] = max(dl, ans[a[k - zam]])
            l += 1
            if l < n:
                if s[l] == '0':
                    zam += 1
            dl += 1
        l = i
        zam = 0
        dl = 0
        while l < n:
            if zam <= k:
                ans[dl] = max(ans[dl], b[k - zam])
            l += 1
            if l < n:
                if s[l] == '1':
                    zam += 1
            dl += 1
    for A in range(1, n + 1):
        ANSWER = 0
        for j in range(0, n + 1):
            if ans[j] >= 0:
                ANSWER = max(ANSWER, A * j + ans[j])
        print(ANSWER, end=' ')
    print()
t = int(input())
for i in range(t):
    (a, b) = map(int, input().split())
    s = input()
    main(a, b, s)

Solution (Python):
t = int(input())
for _ in range(t):
    (n, m) = map(int, input().split())
    fl = [[0 for j in range(n + 2)] for i in range(n + 2)]
    fr = [[0 for j in range(n + 2)] for i in range(n + 2)]
    str = input()
    a = [0 for i in range(n + 2)]
    ans = [-1000000000 for i in range(n + 2)]
    for i in range(n):
        a[i + 1] = ord(str[i]) - ord('0')
    for i in range(1, n + 1):
        s = 0
        for j in range(i, n + 1):
            s += a[j] != 1
            fl[j][s] = max(fl[j][s], j - i + 1)
            fr[i][s] = max(fr[i][s], j - i + 1)
    for i in range(1, n + 1):
        for s in range(0, n + 1):
            fl[i][s] = max(fl[i][s], fl[i - 1][s])
            if s:
                fl[i][s] = max(fl[i][s], fl[i][s - 1])
    for i in reversed(range(1, n + 1)):
        for s in range(0, n + 1):
            fr[i][s] = max(fr[i][s], fr[i + 1][s])
            if s:
                fr[i][s] = max(fr[i][s], fr[i][s - 1])
    ans[0] = max(ans[0], fl[n][m])
    for i in range(1, n + 1):
        s = 0
        for j in range(i, n + 1):
            s += a[j] != 0
            if s > m:
                break
            mx1 = max(fl[i - 1][m - s], fr[j + 1][m - s])
            ans[j - i + 1] = max(ans[j - i + 1], mx1)
    for a in range(1, n + 1):
        mx = 0
        for i in range(0, n + 1):
            mx = max(mx, i * a + ans[i])
        print(mx, end=' ')
    print('')

Solution (Python):
import sys
input = sys.stdin.buffer.readline

def process(S, k):
    n = len(S)
    leftmost_zero = [[0 for j in range(n)] for i in range(k + 1)]
    rightmost_zero = [[0 for j in range(n)] for i in range(k + 1)]
    leftmost_one = [[0 for j in range(n)] for i in range(k + 1)]
    rightmost_one = [[0 for j in range(n)] for i in range(k + 1)]
    for i in range(n):
        one_sum = 0
        zero_sum = 0
        for j in range(i, -1, -1):
            if S[j] == '0':
                zero_sum += 1
            else:
                one_sum += 1
            s = zero_sum + one_sum
            if one_sum <= k:
                leftmost_zero[one_sum][i] = max(leftmost_zero[one_sum][i], s)
                rightmost_zero[one_sum][i - s + 1] = max(rightmost_zero[one_sum][i - s + 1], s)
            if zero_sum <= k:
                leftmost_one[zero_sum][i] = max(leftmost_one[zero_sum][i], s)
                rightmost_one[zero_sum][i - s + 1] = max(rightmost_one[zero_sum][i - s + 1], s)
    for i in range(k + 1):
        for j in range(n - 1, -1, -1):
            if i > 0:
                rightmost_zero[i][j] = max(rightmost_zero[i][j], rightmost_zero[i - 1][j])
            if j + 1 <= n - 1:
                rightmost_zero[i][j] = max(rightmost_zero[i][j], rightmost_zero[i][j + 1])
            if i > 0:
                rightmost_one[i][j] = max(rightmost_one[i][j], rightmost_one[i - 1][j])
            if j + 1 <= n - 1:
                rightmost_one[i][j] = max(rightmost_one[i][j], rightmost_one[i][j + 1])
    pairs = [-1 for i in range(n + 1)]
    for i in range(k + 1):
        for r1 in range(n):
            s1 = leftmost_zero[i][r1]
            pairs[s1] = max(pairs[s1], 0)
            if r1 + 1 <= n - 1:
                s2 = rightmost_one[k - i][r1 + 1]
                pairs[s1] = max(pairs[s1], s2)
            s1 = leftmost_one[i][r1]
            pairs[0] = max(pairs[0], s1)
            if r1 + 1 <= n - 1:
                s2 = rightmost_zero[k - i][r1 + 1]
                pairs[s2] = max(pairs[s2], s1)
    answer = [0 for i in range(n)]
    for s1 in range(n + 1):
        s2 = pairs[s1]
        if s2 != -1:
            for a in range(1, n + 1):
                answer[a - 1] = max(answer[a - 1], a * s1 + s2)
    answer = ' '.join(map(str, answer))
    sys.stdout.write(f'{answer}\n')
    return
t = int(input())
for i in range(t):
    (n, k) = [int(x) for x in input().split()]
    S = input().decode().strip()
    process(S, k)