D - 荷物運びの道 / The Path of Carrying Luggage Editorial by admin
Claude 4.6 Opus (Thinking)概要
一列に並んだ倉庫を東へ順に訪れ、累積重量が上限 \(K\) を超えた倉庫で止まる問題です。各スタート地点に対する停止倉庫番号 \(f(s)\) を求め、区間の総和を高速に答えます。
考察
問題の言い換え
倉庫 \(s\) からスタートすると、荷物を \(A_s, A_{s+1}, \ldots\) と順に拾っていきます。累積重量が初めて \(K\) を超える倉庫 \(t\) で止まります。超えずに倉庫 \(N\) まで到達すれば \(f(s) = N\) です。
累積重量は 累積和(prefix sum) を使って表現できます。\(P[i] = A_1 + A_2 + \cdots + A_i\)(\(P[0] = 0\))と定義すると、倉庫 \(s\) から倉庫 \(t\) までの荷物の合計は:
\[A_s + A_{s+1} + \cdots + A_t = P[t] - P[s-1]\]
よって、\(f(s)\) は次のように定義できます:
\(P[t] - P[s-1] > K\)、すなわち \(P[t] > P[s-1] + K\) を満たす 最小の \(t \geq s\) が \(f(s)\)。そのような \(t\) が \(N\) 以下に存在しなければ \(f(s) = N\)。
素朴なアプローチの問題点
各スタート地点 \(s\) に対して、倉庫を一つずつ調べて累積重量が \(K\) を超えるか確認すると、最悪 \(O(N)\) かかります。全スタート地点で \(O(N^2)\)、さらにクエリもあるので間に合いません。
解決の鍵:二分探索
\(A_i \geq 1\) より \(P\) は単調非減少です。したがって、\(P[t] > P[s-1] + K\) を満たす最小の \(t\) は 二分探索 で \(O(\log N)\) で求められます。
クエリへの対応
\(f(s)\) をすべての \(s\) について前計算した後、\(f\) の累積和 \(Sf[i] = f[1] + f[2] + \cdots + f[i]\) を作れば、各クエリ \(\sum_{s=L}^{R} f(s) = Sf[R] - Sf[L-1]\) と \(O(1)\) で答えられます。
アルゴリズム
累積和の構築: \(P[0] = 0\), \(P[i] = P[i-1] + A[i]\) を計算する。
\(f(s)\) の計算(各 \(s\) に対して二分探索):
- \(\text{target} = P[s-1] + K\) とする。
- \(P[N] \leq \text{target}\) ならば、倉庫 \(N\) まで到達できるので \(f(s) = N\)。
- そうでなければ、\(P[t] > \text{target}\) となる最小の \(t\)(\(s \leq t \leq N\))を二分探索で求め、\(f(s) = t\) とする。
具体例: \(N=5\), \(K=7\), \(A = [3, 2, 4, 1, 5]\) の場合、\(P = [0, 3, 5, 9, 10, 15]\)。 - \(s=1\): target \(= 0 + 7 = 7\)。\(P[3]=9 > 7\) で \(P[2]=5 \leq 7\) なので \(f(1) = 3\)。 - \(s=3\): target \(= 5 + 7 = 12\)。\(P[5]=15 > 12\) で \(P[4]=10 \leq 12\) なので \(f(3) = 5\)。
\(f\) の累積和の構築: \(Sf[i] = Sf[i-1] + f[i]\)。
クエリ応答: 各クエリ \((L_j, R_j)\) に対し \(Sf[R_j] - Sf[L_j - 1]\) を出力。
計算量
- 時間計算量: \(O(N \log N + Q)\)
- \(f(s)\) の計算に二分探索 \(O(\log N)\) × \(N\) 回 = \(O(N \log N)\)
- 各クエリは累積和により \(O(1)\)
- 空間計算量: \(O(N)\)
- 累積和配列 \(P\), \(f\), \(Sf\) にそれぞれ \(O(N)\)
実装のポイント
\(K\) が最大 \(10^{15}\)、\(A_i\) の合計も最大 \(2 \times 10^{14}\) 程度になるため、64ビット整数(Python では自動対応)を使う必要があります。
累積和 \(P\) は単調非減少であることが二分探索の前提です。\(A_i \geq 1\) が保証されているので問題ありません。
二分探索の範囲は \([s, N]\) です。探索対象の配列全体ではなく、\(s\) 以降だけを見ることに注意してください。
高速な入出力のため、
sys.stdin.buffer.read()でまとめて読み込み、sys.stdout.write()でまとめて出力しています。ソースコード
import sys
from bisect import bisect_right
def main():
input_data = sys.stdin.buffer.read().split()
idx = 0
N = int(input_data[idx]); idx += 1
K = int(input_data[idx]); idx += 1
Q = int(input_data[idx]); idx += 1
A = [0] * (N + 1)
for i in range(1, N + 1):
A[i] = int(input_data[idx]); idx += 1
# prefix sum: P[0] = 0, P[i] = A[1] + ... + A[i]
P = [0] * (N + 1)
for i in range(1, N + 1):
P[i] = P[i - 1] + A[i]
# f(s): starting from warehouse s, find the warehouse where we stop
# We accumulate A[s] + A[s+1] + ... + A[t] = P[t] - P[s-1]
# We stop at t if P[t] - P[s-1] > K, i.e., P[t] > P[s-1] + K
# If no such t <= N, then f(s) = N
# Otherwise f(s) = smallest t >= s such that P[t] > P[s-1] + K
# Compute f(s) for all s using binary search
f = [0] * (N + 2) # f[1..N]
for s in range(1, N + 1):
target = P[s - 1] + K
# Find smallest t in [s, N] such that P[t] > target
# P is non-decreasing, so we can binary search
# bisect_right finds insertion point for target in P[s..N]
# We want first index where P[t] > target
lo, hi = s, N
result = N # default: reach warehouse N
if P[N] <= target:
f[s] = N
else:
# binary search for smallest t in [s, N] with P[t] > target
lo, hi = s, N
while lo < hi:
mid = (lo + hi) // 2
if P[mid] > target:
hi = mid
else:
lo = mid + 1
f[s] = lo
# Now we need to answer Q queries: sum of f(L) + f(L+1) + ... + f(R)
# Precompute prefix sums of f
Sf = [0] * (N + 2)
for i in range(1, N + 1):
Sf[i] = Sf[i - 1] + f[i]
out = []
for _ in range(Q):
L = int(input_data[idx]); idx += 1
R = int(input_data[idx]); idx += 1
out.append(str(Sf[R] - Sf[L - 1]))
sys.stdout.write('\n'.join(out) + '\n')
main()
この解説は claude4.6opus-thinking によって生成されました。
posted:
last update: