概要

\(N\) 日間の旅行から連続する区間（または \(0\) 日間）を取りやめ、旅行にかかる費用の合計をちょうど \(K\) 円にする方法が何通りあるかを求める問題です。 「取りやめる費用の合計」に着目し、条件を整理した上で尺取り法を用いることで \(O(N)\) で解くことができます。

考察

1. 問題の言い換え

旅行全体の総費用を \(S = \sum_{i=1}^N A_i\) とします。 残る費用をちょうど \(K\) 円にするということは、取りやめる費用の合計をちょうど \(T = S - K\) 円にすることと同値です。

取りやめる費用の選択肢は次の2通りです。 1. 1日も取りやめない場合: 取りやめる費用は \(0\) 円 2. 区間 \([l, r]\) を取りやめる場合: 取りやめる費用は \(\sum_{i=l}^r A_i\) 円

制約より各日の費用は \(A_i \ge 1\) なので、1日以上取りやめる場合の費用は必ず \(1\) 以上になります。この性質から、\(T\) の値によって場合分けができます。

• \(T < 0\) のとき: 取りやめる費用を負にすることはできないため、達成不可能です。答えは 0 通りとなります。
• \(T = 0\) のとき: 取りやめる費用を \(0\) にする必要があるため、「1日も取りやめない」という \(1\) 通りのみとなります（区間を取りやめると \(A_i \ge 1\) より必ず \(1\) 以上減ってしまいます）。答えは 1 通りです。
• \(T > 0\) のとき: 「1日も取りやめない」選択肢は費用が \(0\) なので除外されます。したがって、連続部分列の和がちょうど \(T\) となる区間 \([l, r]\) の個数を求めればよいことになります。

2. 素朴なアプローチと高速化

すべての区間 \([l, r]\) \((1 \le l \le r \le N)\) について総和を計算する方法では、区間の数が \(\frac{N(N+1)}{2}\) 通りあるため \(O(N^2)\) の計算量がかかります。本問の制約は \(N \le 10^6\) であるため、実行時間制限に間に合いません（TLE）。

ここで、すべての要素が正数（\(A_i > 0\)）であることに着目します。 区間の右端を右に伸ばすと和は単調に増加し、左端を右に縮めると和は単調に減少します。この単調性を利用することで、尺取り法（Two Pointers） を用いて \(O(N)\) で探索できます。

アルゴリズム

尺取り法を用いて、連続部分列の和が \(T\) となる区間を数え上げます。

1. 全体の合計 \(S\) を計算し、\(T = S - K\) を求める。
2. \(T < 0\) なら 0 を、\(T = 0\) なら 1 を出力して終了する。
3. 区間の左端 left、右端 right、現在の区間和 cur、条件を満たす区間数 ans を用意する。
4. right を \(0\) から \(N-1\) まで順に進めながら、以下を行う：
   • cur に \(A[\text{right}]\) を加える。
   • cur > T である間、cur から \(A[\text{left}]\) を引き、left を \(1\) 進める。
   • cur == T となった場合、和が \(T\) となる区間が見つかったので ans を \(1\) 増やす。
5. 最終的な ans を出力する。

計算量

• 時間計算量: \(O(N)\) right は \(0\) から \(N-1\) まで \(N\) 回進みます。また、left も高々 \(N\) 回しか進まないため、ループ内の処理は全体で高々 \(2N\) 回しか実行されません。よって \(O(N)\) となり、十分に間に合います。
• 空間計算量: \(O(N)\) 配列 \(A\) を保持するために \(O(N)\) のメモリを使用します。

実装のポイント

• \(N \le 10^6\) と入力サイズが大きいため、Python では sys.stdin.read を用いて一括で入力を取得することで高速化を図っています。

• \(A_i \ge 1\) であるため、区間の総和が \(T\) になるような左端 left は各 right に対して高々 \(1\) つしか存在しません。そのため、cur == T の判定でそのまま ans += 1 とカウントして問題ありません。

  ソースコード

import sys


def main():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return

    N = int(data[0])
    K = int(data[1])
    A = [int(x) for x in data[2:]]

    S = sum(A)
    T = S - K

    if T < 0:
        print(0)
        return
    if T == 0:
        print(1)
        return

    ans = 0
    cur = 0
    left = 0
    for right in range(N):
        cur += A[right]
        while cur > T:
            cur -= A[left]
            left += 1
        if cur == T:
            ans += 1

    print(ans)


if __name__ == "__main__":
    main()

---

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