トップ > AtCoder > ABC216F - Max Sum Counting

ABC216F - Max Sum Counting

AtCoder AtCoder-Beginner Contest DP 数え上げ 数え上げ-不等式 Python

考え方

  • 「dp[i][j] = $i$番目まで選んだときに和が$j$となる場合の数」としたくなる
  • 不等式の両辺が動くと困るので，片方を固定したい．和については↑のようにdpで考えたいので，maxを固定したい．

$\{(A_{i}, B_{i})\}_{i}$を$A_{i}$についてソートして昇順に考えれば，「$i$番目まで選ぶ & $i$番目を必ず選ぶ」とき「$\max = A_{i}$」となる．dpを計算する各ステップで$\max$が固定されていれば，各ステップで適切な$j$の範囲のdpを答えに加算していけばよさそう．


$i$まで計算できているとき，$i+1$は

\begin{aligned}
\mathrm{dp}[i+1][j+b]
&= \underbrace{\mathrm{dp}[i][j]}_{B_{i}\text{を選ぶ場合}}
+ \underbrace{\mathrm{dp}[i][j+b]}_{B_{i}\text{を選ばない場合}}
\end{aligned}
で求められる（dp[i]は$i$を選んでも選ばなくてもよい）．


dp[i]がわかれば，「$i$番目まで選ぶ & $i$番目を必ず選ぶ」ときの答えは

\begin{aligned}
\sum_{j=0}^{A_{i} - B_{i}} \mathrm{dp}[i - 1][j]
\end{aligned}
で計算できる（$i$番目を必ず選ぶことは，$B_{i}$を予め引くことで表現できる）．

よって，最終的に求めたいのは，すべての$i$の寄与を足し合わせた

\begin{aligned}
\sum_{\substack{i \\ (A_{i} \geq B_{i}) }}
\sum_{j=0}^{A_{i} - B_{i}} \mathrm{dp}[i - 1][j]
\pmod M
\end{aligned}
となる．

回答例

dp[i + 1] = dp[i]とすると，同じリストを参照することになりバグる．

$j$の範囲は$\max A_{i}$まであれば十分．

N = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
AB = sorted([(a, b) for a, b in zip(A, B)])
mod = 998244353

dp = [[0] * 5001 for _ in range(N + 1)]
dp[0][0] = 1

ans = 0
for i in range(N):
    a, b = AB[i]
    for j in range(5001):
        dp[i + 1][j] += dp[i][j]
        dp[i + 1][j] %= mod
        if j + b <= 5000:
            dp[i + 1][j + b] += dp[i][j]
            dp[i + 1][j + b] %= mod
        if j + b <= a:
            ans += dp[i][j]
            ans %= mod
print(ans)

【参考】Submission #25432523 - AtCoder Beginner Contest 216

参考