Hello, I am currently doing this question: https://mirror.codeforces.com/contest/837/problem/D

And I need some help. I came up with a good dp algorithm myself:

dp[i][j][k] means the maximum exponent of prime factor 2 we can achieve when we:

1. choose from the first i integers in a

2. choose a subset of length j from the first i integers

3. the exponent of 5 is at most k

then, dp[i][j][k] = max(dp[i-1][j][k], dp[i-1][j-1][k-cur_exp5] + cur_exp2)

note dp[0][0][0] = 0 (don't have to initialize anything)

And implemented this algorithm in Python:

n, K = map(int, input().split()) a = list(map(int, input().split()))

def factors(x): cnt2, cnt5 = 0, 0 while x % 5 == 0: x //= 5 cnt5 += 1 while x % 2 == 0: x //= 2 cnt2 += 1 return [cnt2, cnt5]

b = [] mc = [] for A in a: c2, c5 = factors(A) b.append([c2, c5]) mc.append(c5) mc.sort(reverse=True) MAX_CNT5 = sum(mc[:K])

dp = [[[0 for _ in range(MAX_CNT5+1)] for _ in range(K+1)] for _ in range(n+1)]

for i in range(1, n+1): for j in range(1, K+1): for k in range(MAX_CNT5+1): if k >= b[i-1][1]: dp[i][j][k] = max(dp[i-1][j][k], dp[i-1][j-1][k-b[i-1][1]] + b[i-1][0]) else: dp[i][j][k] = dp[i-1][j][k]

ans = 0 for i in range(len(dp[n][K])): ans = max(ans, min(i, dp[n][K][i]))

print(ans)

However, this code WAs for test case 8, with my answer of 11 being higher than the expected answer of 9. I have read the editorial for this question, and I believe the DP formula in the editorial is the same as my DP formula mentioned above (please correct me if I am wrong). Therefore, because I believe I have the DP formula correct, I am really confused on why this code still WAs. I know asking people to debug a code is sometimes annoying but I really tried everything and could not figure out why this code is not working, and I would really appreciate it if anyone could tell me why.

Thank you so much!