D - 均等な買い物 / Equal Shopping Editorial by admin
gpt-5.5-high概要
各お店について「買わない」「高橋君が買う」「青木君が買う」のいずれかを選び、高橋君と青木君の購入金額の合計が等しくなる計画数を数える問題です。
\(P,Q \leq 3\)、購入金額は最大 \(20\) と小さいため、「何軒選んだか」と「合計金額の差」を状態にした DP で数えます。
考察
高橋君の合計金額を \(S_T\)、青木君の合計金額を \(S_A\) とします。
均等な買い物計画では、
\[ S_T = S_A \]
である必要があります。
これは差
\[ S_T - S_A \]
が \(0\) である、と言い換えられます。
そこで、各お店を順番に見ながら、以下を管理する DP を考えます。
- 高橋君がすでに買ったお店の数
- 青木君がすでに買ったお店の数
- 現在の金額差 \(S_T - S_A\)
各お店では次の 3 通りがあります。
- 誰も買わない
→ 状態は変わらない - 高橋君が金額 \(x\) で買う
→ 高橋君の店数が \(1\) 増え、差が \(+x\) される - 青木君が金額 \(x\) で買う
→ 青木君の店数が \(1\) 増え、差が \(-x\) される
最後に、
- 高橋君が \(P\) 軒
- 青木君が \(Q\) 軒
- 差が \(0\)
である状態の数が答えになります。
差の範囲
高橋君が買う金額は最大で \(20P\)、青木君が買う金額は最大で \(20Q\) です。
したがって差 \(S_T - S_A\) の範囲は
\[ -20Q \leq S_T - S_A \leq 20P \]
です。
負の添字は配列で扱いづらいので、コードでは
\[ \text{OFF} = 20Q \]
だけずらして管理しています。
つまり、差 \(d\) は配列上では
\[ d + \text{OFF} \]
番目に保存します。
素朴な方法が難しい理由
各お店について 3 通りの状態があり、さらに買う場合は最大 20 通りの金額があります。
全通りを列挙すると、指数時間になってしまい、\(N\) が最大 \(2000\) では到底間に合いません。
一方で、\(P,Q \leq 3\)、金額も最大 \(20\) なので、必要な状態数は非常に小さいです。
そのため、DP によって効率よく数えることができます。
また、更新は \(M\) 回ありますが、制約で
\[ N \times M \leq 2000 \]
が保証されています。
したがって、各更新のたびに DP を最初から計算し直しても十分間に合います。
アルゴリズム
DP を次のように定義します。
\[ dp[p][q][d] \]
を、
- 高橋君が \(p\) 軒で買った
- 青木君が \(q\) 軒で買った
- 差 \(S_T - S_A\) が \(d\)
となる買い物計画の個数とします。
実装では \(d\) をそのまま添字にできないため、実際には
\[ dp[p][q][d + \text{OFF}] \]
として持ちます。
初期状態は、まだ誰も何も買っていないので、
\[ dp[0][0][0] = 1 \]
です。
各お店の購入可能金額が \([L,R]\) のとき、遷移は以下の通りです。
1. 誰も買わない
状態は変わりません。
\[ dp[p][q][d] \to dp[p][q][d] \]
これはコードでは ndp = [arr[:] for arr in dp] として、現在の DP をそのままコピーすることで実現しています。
2. 高橋君が買う
高橋君が金額 \(x\) で買うと、差が \(+x\) されます。
\[ dp[p][q][d] \to dp[p+1][q][d+x] \]
ただし、\(p < P\) の場合のみです。
3. 青木君が買う
青木君が金額 \(x\) で買うと、差が \(-x\) されます。
\[ dp[p][q][d] \to dp[p][q+1][d-x] \]
ただし、\(q < Q\) の場合のみです。
区間和による高速化
購入金額 \(x\) は \(L \leq x \leq R\) の範囲です。
単純にすべての \(x\) を試しても制約上はかなり小さいですが、コードでは累積和を使ってまとめて遷移しています。
例えば高橋君が買う場合、遷移先の差を \(j\) とすると、遷移元は
\[ j-R, j-R+1, \dots, j-L \]
の範囲です。
この範囲の DP 値の合計を累積和で \(O(1)\) で求めています。
青木君が買う場合も同様に、遷移先の差を \(j\) とすると、遷移元は
\[ j+L, j+L+1, \dots, j+R \]
です。
全てのお店を処理した後、求める答えは
\[ dp[P][Q][0] \]
です。
コード上では差 \(0\) は添字 OFF に対応するので、
dp[TARGET][OFF]
を出力しています。
ここで TARGET は \((P,Q)\) に対応する添字です。
計算量
差の取りうる範囲の長さを
\[ D = 20(P+Q)+1 \]
とします。
状態数は
\[ (P+1)(Q+1)D \]
です。
各更新ごとに全 \(N\) 店を DP で処理するので、
- 時間計算量: \(O\left(MN(P+1)(Q+1)D\right)\)
- 空間計算量: \(O\left((P+1)(Q+1)D\right)\)
です。
制約より \(P,Q \leq 3\)、\(D \leq 121\)、さらに \(NM \leq 2000\) なので十分高速です。
実装のポイント
- 差 \(S_T-S_A\) は負になる可能性があるため、
OFF = 20 * Qだけずらして配列で管理します。 - DP の添字を簡単にするため、コードでは \((p,q)\) を一次元にまとめています。
index = p * (Q + 1) + q
各お店について、まず「誰も買わない」場合を反映するために現在の DP をコピーして
ndpを作ります。高橋君・青木君が買う遷移では、購入金額の範囲 \([L,R]\) に対する和を累積和で求めています。
答えは、更新を反映した直後に毎回 DP を計算し直して求めます。制約で \(NM \leq 2000\) なので、この方針で間に合います。
ソースコード
import sys
MOD = 998244353
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
N = next(it)
P = next(it)
Q = next(it)
M = next(it)
ranges = []
for _ in range(N):
l = next(it)
r = next(it)
ranges.append((l, r))
Q1 = Q + 1
P1 = P + 1
SZ = P1 * Q1
OFF = 20 * Q
D = 20 * (P + Q) + 1
TARGET = P * Q1 + Q
def calc():
dp = [[0] * D for _ in range(SZ)]
dp[0][OFF] = 1
for l, r in ranges:
ndp = [arr[:] for arr in dp]
for p in range(P1):
base = p * Q1
for q in range(Q1):
if p == P and q == Q:
continue
arr = dp[base + q]
pref = [0] * (D + 1)
s = 0
for i in range(D):
s += arr[i]
pref[i + 1] = s
if s == 0:
continue
if p < P:
dest = ndp[base + Q1 + q]
for j in range(l, D):
left = j - r
if left < 0:
left = 0
dest[j] += pref[j - l + 1] - pref[left]
if q < Q:
dest = ndp[base + q + 1]
end = D - l
for j in range(end):
right = j + r + 1
if right > D:
right = D
dest[j] += pref[right] - pref[j + l]
for arr in ndp:
for i in range(D):
if arr[i] >= MOD:
arr[i] %= MOD
dp = ndp
return dp[TARGET][OFF] % MOD
ans = []
for _ in range(M):
x = next(it) - 1
a = next(it)
b = next(it)
ranges[x] = (a, b)
ans.append(str(calc()))
sys.stdout.write("\n".join(ans))
if __name__ == "__main__":
main()
この解説は gpt-5.5-high によって生成されました。
posted:
last update: