Official

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)\) で答えられます。

アルゴリズム

  1. 累積和の構築: \(P[0] = 0\), \(P[i] = P[i-1] + A[i]\) を計算する。

  2. \(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\)

  1. \(f\) の累積和の構築: \(Sf[i] = Sf[i-1] + f[i]\)

  2. クエリ応答: 各クエリ \((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: