Official

E - お花見の予算チェック / Cherry Blossom Viewing Budget Check Editorial by admin

gemini-3-flash-thinking

概要

この問題は、長さ \(N\) の数列の中から指定された区間 \([L, R]\) の総和を求め、それが予算 \(T\) を超えるかどうかを \(Q\) 個のクエリに対して判定する問題です。

考察

まず、素直に各クエリに対して「\(L\) 番目から \(R\) 番目までのお店をループで回して合計金額を計算する」という方法を考えてみます。 この場合、1つのクエリあたり最大で \(N\) 回の足し算が必要になります。クエリの数は \(Q\) 個あるため、全体の計算量は \(O(N \times Q)\) となります。

制約を確認すると、\(N\) と \(Q\) はそれぞれ最大 \(2 \times 10^5\) です。最悪の場合 \(2 \times 10^5 \times 2 \times 10^5 = 4 \times 10^{10}\) 回程度の計算が必要になり、一般的な実行時間制限(2秒程度)には間に合いません(TLE:Time Limit Exceeded となります)。

そこで、「区間の和を高速に求める手法」が必要になります。これには 累積和(Prefix Sum) というアルゴリズムが最適です。

アルゴリズム

累積和の利用

累積和とは、配列の先頭からある要素までの合計値をあらかじめ計算しておく手法です。

  1. 長さ \(N+1\) の配列 \(S\) を用意します。
  2. \(S[0] = 0\) とし、各 \(i\) (\(1 \leq i \leq N\)) について \(S[i] = A_1 + A_2 + \cdots + A_i\) を計算します。これは \(S[i] = S[i-1] + A_i\) という式で、前から順番に足していくことで \(O(N)\) で作成できます。
  3. \(L\) 番目から \(R\) 番目までの合計金額を求めたいときは、以下の式で計算できます。 $\(\text{合計金額} = S[R] - S[L-1]\)$

たとえば、1番目から3番目までの和が欲しい場合、\(S[3] = (A_1 + A_2 + A_3)\) から \(S[0] = 0\) を引けば求められます。2番目から3番目までの和なら、\(S[3] = (A_1 + A_2 + A_3)\) から \(S[1] = (A_1)\) を引くことで \((A_2 + A_3)\) が得られます。

この工夫により、各クエリに対して引き算1回(\(O(1)\))で答えを出せるようになります。

計算量

  • 時間計算量: \(O(N + Q)\)
    • 累積和の構築に \(O(N)\)、各クエリの処理に \(O(1)\) かかるため、全体で \(O(N + Q)\) となります。今回の制約(\(2 \times 10^5\))では、十分高速に動作します。
  • 空間計算量: \(O(N)\)
    • 長さ \(N+1\) の累積和配列 \(S\) を保持するために必要なメモリ量です。

実装のポイント

  • 高速な入出力: Pythonでは入力数や出力数が多い場合、input() を繰り返すと時間がかかることがあります。sys.stdin.read().split() で一括で読み込み、sys.stdout.write() でまとめて出力することで実行時間を大幅に短縮できます。

  • インデックスの扱い: 問題文は 1-indexed(1番目から始まる)ですが、プログラム内での配列の扱いや累積和の定義(\(S[0]=0\) とする慣習)と整合性をとる必要があります。累積和配列を \(N+1\) のサイズで作っておくと、1-indexed の \(L, R\) をそのまま活用しやすくなります。

  • 大きな数値: 予算上限 \(T\) や合計金額は最大で \(10^{14}\) 程度になります。Pythonは整数の大きさに制限がないためそのまま計算できますが、他の言語(C++など)では 64bit 整数型(long long など)を使用する必要があります。

    ソースコード

import sys

def solve():
    # 入力を一度に読み込む(高速化のため)
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # 基本情報の取得
    N = int(input_data[0])
    Q = int(input_data[1])
    T = int(input_data[2])
    
    # 各お店の価格を取得
    A = list(map(int, input_data[3:3+N]))
    
    # 累積和の計算
    # S[i] は A[0] から A[i-1] までの合計
    S = [0] * (N + 1)
    for i in range(N):
        S[i+1] = S[i] + A[i]
        
    # クエリの処理
    queries = input_data[3+N:]
    results = []
    for i in range(0, 2 * Q, 2):
        L = int(queries[i])
        R = int(queries[i+1])
        
        # L番目からR番目までの合計金額を累積和で計算 (1-indexed)
        # S[R] = A[0] + ... + A[R-1]
        # S[L-1] = A[0] + ... + A[L-2]
        # S[R] - S[L-1] = A[L-1] + ... + A[R-1]
        total_price = S[R] - S[L-1]
        
        if total_price > T:
            results.append("Yes")
        else:
            results.append("No")
            
    # 結果を一括で出力
    sys.stdout.write("\n".join(results) + "\n")

if __name__ == "__main__":
    solve()

この解説は gemini-3-flash-thinking によって生成されました。

posted:
last update: