Boxes and Balls (Codeforces 1845E)
Tags: Math, dp

There are $$$n$$$ boxes placed in a line. The boxes are numbered from $$$1$$$ to $$$n$$$. Some boxes contain one ball inside of them, the rest are empty. At least one box contains a ball and at least one box is empty.In one move, you have to choose a box with a ball inside and an adjacent empty box and move the ball from one box into another. Boxes $$$i$$$ and $$$i+1$$$ for all $$$i$$$ from $$$1$$$ to $$$n-1$$$ are considered adjacent to each other. Boxes $$$1$$$ and $$$n$$$ are not adjacent.How many different arrangements of balls exist after exactly $$$k$$$ moves are performed? Two arrangements are considered different if there is at least one such box that it contains a ball in one of them and doesn't contain a ball in the other one.Since the answer might be pretty large, print its remainder modulo $$$10^9+7$$$.

Input: The first line contains two integers $$$n$$$ and $$$k$$$ ($$$2 \le n \le 1500$$$; $$$1 \le k \le 1500$$$) — the number of boxes and the number of moves.The second line contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$a_i \in \{0, 1\}$$$) — $$$0$$$ denotes an empty box and $$$1$$$ denotes a box with a ball inside. There is at least one $$$0$$$ and at least one $$$1$$$.

Output: Print a single integer — the number of different arrangements of balls that can exist after exactly $$$k$$$ moves are performed, modulo $$$10^9+7$$$.

Note: In the first example, there are the following possible arrangements:   0 1 1 0 — obtained after moving the ball from box $$$1$$$ to box $$$2$$$;  1 0 0 1 — obtained after moving the ball from box $$$3$$$ to box $$$4$$$;  1 1 0 0 — obtained after moving the ball from box $$$3$$$ to box $$$2$$$. In the second example, there are the following possible arrangements:   1 0 1 0 — three ways to obtain that: just reverse the operation performed during the first move;  0 1 0 1 — obtained from either of the first two arrangements after the first move. 

Sample Inputs:
4 1
1 0 1 0

4 2
1 0 1 0

10 6
1 0 0 1 0 0 0 1 1 1

Sample Outputs:
3

2

69

Solution (Python):
from bisect import bisect, bisect_left
import sys
input = lambda : sys.stdin.readline().rstrip()
sint = lambda : int(input())
mint = lambda : map(int, input().split())
ints = lambda : list(map(int, input().split()))
MOD = 10 ** 9 + 7

def solve() -> None:
    (n, k) = mint()
    nums = ints()
    less = sum(nums) <= n // 2
    pos = [i for (i, x) in enumerate(nums) if x == less]
    dp = [[0] * (k + 1) for _ in range(len(pos) + 1)]
    dp[0][0] = 1
    for i in range(n):
        j0 = bisect_left(pos, i)
        for j in range(min(i + 1, len(pos), j0 + 40), max(0, j0 - 40), -1):
            d = abs(pos[j - 1] - i)
            for t in range(k, d - 1, -1):
                dp[j][t] = (dp[j][t] + dp[j - 1][t - d]) % MOD
    print(sum((dp[-1][i] for i in range(k, -1, -2))) % MOD)
solve()

Solution (Python):
import sys
input = sys.stdin.readline
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
index = []
for i in range(n):
    if a[i] == 1:
        index.append(i)
MOD = 10 ** 9 + 7
m = len(index)
dp = [[0] * 1501 for _ in range(1501)]
for i in range(k + 1):
    dp[0][i] = int(i % 2 == 0)
j0 = 0
for i in range(1, n + 1):
    if j0 < m and index[j0] < i:
        j0 += 1
    for j in range(min(i, m, j0 + 40) - 1, max(0, j0 - 40) - 1, -1):
        dis = abs(index[j] - i + 1)
        for t in range(k, dis - 1, -1):
            dp[j + 1][t] += dp[j][t - dis]
            dp[j + 1][t] %= MOD
print(dp[m][k])

Solution (Python):
import sys
input = sys.stdin.readline
(n, k) = map(int, input().split())
a = list(map(int, input().split()))
index = []
for i in range(n):
    if a[i] == 1:
        index.append(i)
MOD = 10 ** 9 + 7
m = len(index)
dp = [[0] * 1501 for _ in range(1501)]
for i in range(k + 1):
    dp[0][i] = int(i % 2 == 0)
j0 = 0
for i in range(1, n + 1):
    if j0 < m and index[j0] < i:
        j0 += 1
    for j in range(min(i, m, j0 + 40) - 1, max(0, j0 - 40) - 1, -1):
        dis = abs(index[j] - i + 1)
        for t in range(k, dis - 1, -1):
            dp[j + 1][t] += dp[j][t - dis]
            dp[j + 1][t] %= MOD
print(dp[m][k])