Miriany and Matchstick (Codeforces 1852D)
Tags: Greedy, dp

Miriany's matchstick is a $$$2 \times n$$$ grid that needs to be filled with characters A or B. He has already filled in the first row of the grid and would like you to fill in the second row. You must do so in a way such that the number of adjacent pairs of cells with different characters$$$^\dagger$$$ is equal to $$$k$$$. If it is impossible, report so.$$$^\dagger$$$ An adjacent pair of cells with different characters is a pair of cells $$$(r_1, c_1)$$$ and $$$(r_2, c_2)$$$ ($$$1 \le r_1, r_2 \le 2$$$, $$$1 \le c_1, c_2 \le n$$$) such that $$$|r_1 - r_2| + |c_1 - c_2| = 1$$$ and the characters in $$$(r_1, c_1)$$$ and $$$(r_2, c_2)$$$ are different.

Input: The first line consists of an integer $$$t$$$, the number of test cases ($$$1 \leq t \leq 1000$$$). The description of the test cases follows.The first line of each test case has two integers, $$$n$$$ and $$$k$$$ ($$$1 \leq n \leq 2 \cdot 10^5, 0 \leq k \leq 3 \cdot n$$$) – the number of columns of the matchstick, and the number of adjacent pairs of cells with different characters required.The following line contains string $$$s$$$ of $$$n$$$ characters ($$$s_i$$$ is either A or B) – Miriany's top row of the matchstick.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2 \cdot 10^5$$$.

Output: For each test case, if there is no way to fill the second row with the number of adjacent pairs of cells with different characters equals $$$k$$$, output "NO". Otherwise, output "YES". Then, print $$$n$$$ characters that a valid bottom row for Miriany's matchstick consists of. If there are several answers, output any of them.

Note: In the first test case, it can be proved that there exists no possible way to fill in row $$$2$$$ of the grid such that $$$k = 1$$$. For the second test case, BABB is one possible answer.The grid below is the result of filling in BABB as the second row. $$$\begin{array}{|c|c|} \hline A & A & A & A \cr \hline B & A & B & B \cr \hline \end{array}$$$ The pairs of different characters are shown below in red: $$$\begin{array}{|c|c|} \hline \color{red}{A} & A & A & A \cr \hline \color{red}{B} & A & B & B \cr \hline \end{array}$$$—————————————————$$$\begin{array}{|c|c|} \hline A & A & \color{red}{A} & A \cr \hline B & A & \color{red}{B} & B \cr \hline \end{array}$$$—————————————————$$$\begin{array}{|c|c|} \hline A & A & A & \color{red}{A} \cr \hline B & A & B & \color{red}{B} \cr \hline \end{array}$$$—————————————————$$$\begin{array}{|c|c|} \hline A & A & A & A \cr \hline \color{red}{B} & \color{red}{A} & B & B \cr \hline \end{array}$$$—————————————————$$$\begin{array}{|c|c|} \hline A & A & A & A \cr \hline B & \color{red}{A} & \color{red}{B} & B \cr \hline \end{array}$$$There are a total of $$$5$$$ pairs, which satisfies $$$k$$$.

Sample Inputs:
4
10 1
ABBAAABBAA
4 5
AAAA
9 17
BAAABBAAB
4 9
ABAB

Sample Outputs:
NO
YES
BABB
YES
ABABAABAB
NO

Solution (Python):
def solve():
    (n, k) = map(int, input().split())
    s = input()
    for i in range(1, n):
        k -= s[i] != s[i - 1]
    if k < 0:
        print('NO')
        return
    dp = [[[] for _ in range(n)] for _ in range(2)]
    dp[0][0].append((s[0] != 'A', s[0] != 'A'))
    dp[1][0].append((s[0] != 'B', s[0] != 'B'))
    for i in range(1, n):
        for t in range(2):
            for (l, r) in dp[t][i - 1]:
                dp[t][i].append((l, r))
                dp[not t][i].append((l + 1, r + 1))
        for t in range(2):
            f = []
            dp[t][i].sort()
            d = s[i] != 'A' if t == 0 else s[i] != 'B'
            for (l, r) in dp[t][i]:
                l += d
                r += d
                if f and l <= f[-1][1] + 1:
                    f[-1] = (f[-1][0], max(f[-1][1], r))
                else:
                    f.append((l, r))
            dp[t][i] = f

    def get(f, x):
        for (l, r) in f:
            if l <= x <= r:
                return True
        return False
    for t in range(2):
        if get(dp[t][n - 1], k):
            print('YES')
            ans = list(s)
            for i in range(n - 1, 0, -1):
                k -= s[i] != 'A' if t == 0 else s[i] != 'B'
                ans[i] = 'A' if t == 0 else 'B'
                if not get(dp[t][i - 1], k):
                    t = not t
                    k -= 1
            ans[0] = 'A' if t == 0 else 'B'
            print(''.join(ans))
            return
    print('NO')
t = int(input())
for _ in range(t):
    solve()

Solution (Python):
def solve():
    (n, k) = map(int, input().split())
    s = input()
    for i in range(1, n):
        k -= s[i] != s[i - 1]
    if k < 0:
        print('NO')
        return
    dp = [[[] for _ in range(n)] for _ in range(2)]
    dp[0][0].append((s[0] != 'A', s[0] != 'A'))
    dp[1][0].append((s[0] != 'B', s[0] != 'B'))
    for i in range(1, n):
        for t in range(2):
            for (l, r) in dp[t][i - 1]:
                dp[t][i].append((l, r))
                dp[not t][i].append((l + 1, r + 1))
        for t in range(2):
            f = []
            dp[t][i].sort()
            d = s[i] != 'A' if t == 0 else s[i] != 'B'
            for (l, r) in dp[t][i]:
                l += d
                r += d
                if f and l <= f[-1][1] + 1:
                    f[-1] = (f[-1][0], max(f[-1][1], r))
                else:
                    f.append((l, r))
            dp[t][i] = f

    def get(f, x):
        for (l, r) in f:
            if l <= x <= r:
                return True
        return False
    for t in range(2):
        if get(dp[t][n - 1], k):
            print('YES')
            ans = list(s)
            for i in range(n - 1, 0, -1):
                k -= s[i] != 'A' if t == 0 else s[i] != 'B'
                ans[i] = 'A' if t == 0 else 'B'
                if not get(dp[t][i - 1], k):
                    t = not t
                    k -= 1
            ans[0] = 'A' if t == 0 else 'B'
            print(''.join(ans))
            return
    print('NO')
t = int(input())
for _ in range(t):
    solve()

Solution (Python):
def solve():
    (n, v) = map(int, input().split())
    s = input()
    for i in range(1, n):
        v -= s[i] != s[i - 1]
    if v < 0:
        print('NO')
        return
    dp = [[[] for _ in range(n)] for _ in range(2)]
    dp[0][0].append((s[0] != 'A', s[0] != 'A'))
    dp[1][0].append((s[0] != 'B', s[0] != 'B'))
    for i in range(1, n):
        for t in range(2):
            for (l, r) in dp[t][i - 1]:
                dp[t][i].append((l, r))
                dp[not t][i].append((l + 1, r + 1))
        for t in range(2):
            f = []
            dp[t][i].sort()
            d = s[i] != 'A' if t == 0 else s[i] != 'B'
            for (l, r) in dp[t][i]:
                l += d
                r += d
                if f and l <= f[-1][1] + 1:
                    f[-1] = (f[-1][0], max(f[-1][1], r))
                else:
                    f.append((l, r))
            dp[t][i] = f

    def get(f, x):
        for (l, r) in f:
            if l <= x <= r:
                return True
        return False
    for t in range(2):
        if get(dp[t][n - 1], v):
            print('YES')
            ans = list(s)
            for i in range(n - 1, 0, -1):
                v -= s[i] != 'A' if t == 0 else s[i] != 'B'
                ans[i] = 'A' if t == 0 else 'B'
                if not get(dp[t][i - 1], v):
                    t = not t
                    v -= 1
            ans[0] = 'A' if t == 0 else 'B'
            print(''.join(ans))
            return
    print('NO')
t = int(input())
for _ in range(t):
    solve()