Official

C - 旅行の計画 / Travel Plan Editorial by admin

Claude Opus 5 (High)

概要

長さ \(N\) の正整数列 \(A\) について、「何も取り除かない」または「連続する区間 \([l, r]\) を 1 つ取り除く」という選び方のうち、残った要素の総和がちょうど \(K\) になるものの個数を数える問題です。

考察

言い換え

全体の総和を \(S = \sum_{i=1}^N A_i\) とおきます。

  • 1 日も取りやめない場合、費用は \(S\) 円。これが \(K\) に等しいかどうかを判定すればよいので、\(S = K\) なら 1 通り、そうでなければ 0 通りです。
  • 区間 \([l, r]\) を取りやめる場合、費用は \(S - \sum_{i=l}^r A_i\) 円です。これが \(K\) になる条件は $\(\sum_{i=l}^{r} A_i = S - K\)$ です。

つまり \(T = S - K\) とおくと、「区間和がちょうど \(T\) になる区間 \((l, r)\) の個数」を求める問題に帰着します。これは非常に典型的な問題です。

なお、\(T < 0\) のときはどんな区間を選んでも達成できません(\(A_i \ge 1\) より区間和は必ず正)。また \(T = 0\) のときも、\(A_i \ge 1\) なので区間和が \(0\) になることはなく、区間を取りやめる方法は 0 通りです(この場合は \(S = K\) なので「1 日も取りやめない」1 通りだけが答えになります)。

素朴なアプローチの問題点

すべての \((l, r)\) の組を全探索すると \(O(N^2)\) 通りあり、\(N \le 10^6\) では \(10^{12}\) 回程度の計算になって間に合いません。

累積和 \(P_i = A_1 + \dots + A_i\) を使い、\(P_r - P_{l-1} = T\) となる組を連想配列(unordered_map など)で数える方法もあります。これは \(O(N)\) 程度ですが、\(N = 10^6\) でハッシュテーブルを使うと定数倍が重く、メモリも余分に必要です。

解決方法:しゃくとり法

ここで重要な観察は、\(A_i \ge 1\)(すべて正)という制約です。これにより、左端 \(l\) を固定したとき区間和 \(\sum_{i=l}^{r} A_i\) は \(r\) について単調増加し、右端 \(r\) を固定したとき \(l\) を大きくすると区間和は単調減少します。

したがって「区間和が \(T\) 以下となる最小の \(l\)」は \(r\) が増えるにつれて単調に増加します。この単調性を利用すると、しゃくとり法(two pointers / 尺取り法)で線形時間に数え上げられます。

アルゴリズム

  1. 入力を読み、総和 \(S\) を計算する。
  2. \(S = K\) なら答えを 1 で初期化(「1 日も取りやめない」方法)、そうでなければ 0 で初期化。
  3. \(T = S - K\) とする。\(T > 0\) のとき、以下のしゃくとり法で「区間和が \(T\) の区間の個数」を数え、答えに加算する。
    • cur(現在の窓 \([l, r]\) の和)を 0、l を先頭にセット。
    • r を \(0\) から \(N-1\) まで動かす:
      • cur += a[r]
      • cur > T の間、cur -= a[l]; l++(窓を左から縮める)
      • この時点で cur <= T。もし cur == T なら、右端が r である条件を満たす区間がちょうど 1 つ見つかったので答えを 1 増やす。
  4. 答えを出力する。

正当性の補足

各 \(r\) に対して、和が \(T\) になる左端 \(l\) は(すべて正なので和が \(l\) について狭義単調減少するため)高々 1 つです。ループ後の l は「\([l,r]\) の和が \(T\) 以下となる最小の \(l\)」なので、和が \(T\) になる \(l\) が存在すればそれはこの l と一致します。よって各 \(r\) について 1 回のチェックで十分です。

具体例

\(N = 5\), \(A = (1, 2, 3, 4, 5)\), \(K = 9\) の場合。\(S = 15\), \(T = 6\)。区間和が 6 になるのは \([1,3]\)(\(1+2+3\))と \([3,3]\)…ではなく \(A_3 = 3\) なので違い、\([2,3]\) は \(5\)、\([3,4]\) は \(7\)。実際には \((l,r) = (1,3)\) のみで 1 通り。\(S \ne K\) なので答えは 1 です。

計算量

  • 時間計算量: \(O(N)\)(r は \(N\) 回、l も合計で高々 \(N\) 回しか進まないため)
  • 空間計算量: \(O(N)\)(数列の保持。読み込みながら処理すれば \(O(1)\) も可能)

実装のポイント

  • オーバーフローに注意: \(S\) は最大で \(10^6 \times 10^9 = 10^{15}\)、\(K\) も最大 \(10^{15}\) なので、総和・累積和・cur はすべて 64 bit 整数(long long)で扱います。int だと確実にオーバーフローします。

  • 「1 日も取りやめない」ケースを忘れない: これは区間を取りやめる場合とは別のカウントです。\(T = 0\)(\(S = K\))のとき、\(A_i \ge 1\) より和が 0 の区間は存在しないので、区間の数え上げは行わなくてよく、答えは 1 になります。

  • \(T < 0\) の場合: 答えは 0(もしくは \(S=K\) 判定の結果のみ)です。\(T > 0\) のときのみループを回すようにガードしておくと安全です。

  • 入出力の高速化: \(N \le 10^6\) なので、C++ では scanf/printf を使うか、cin を使うなら ios::sync_with_stdio(false); cin.tie(nullptr); を入れましょう。Python なら sys.stdin.buffer.read() で一括読み込みするのが必須です。

  • しゃくとり法の内側の while では l <= r の条件も入れておくと、cur が単一要素より大きい場合などでも安全にインデックス範囲を保てます(実際には \(A_i \ge 1, T \ge 1\) なので l > r になると cur = 0 \le T で止まります)。

    ソースコード

#include <bits/stdc++.h>
using namespace std;

int main(){
    int n; long long K;
    scanf("%d %lld", &n, &K);
    vector<long long> a(n);
    long long S=0;
    for(int i=0;i<n;i++){ scanf("%lld",&a[i]); S+=a[i]; }
    long long ans = (S==K)?1:0;
    long long T = S - K;
    if(T > 0){
        long long cur = 0;
        int l = 0;
        for(int r=0;r<n;r++){
            cur += a[r];
            while(cur > T && l <= r){ cur -= a[l]; l++; }
            if(cur == T) ans++;
        }
    }
    printf("%lld\n", ans);
    return 0;
}

この解説は claude5opus-high によって生成されました。