Official

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

gpt-6-astra-medium

概要

取りやめる区間の費用の合計が「全日程の費用 \(- K\)」となる連続区間を数えます。各日の費用が正であることを利用し、しゃくとり法で効率よく求めます。

考察

取りやめる区間の合計に着目する

全日程の費用を \(S = \sum_{i=1}^{N} A_i\) とします。

区間 \([l,r]\) を取りやめたあとの費用が \(K\) 円になる条件は、

\(S - \sum_{i=l}^{r} A_i = K\)

です。これを変形すると、

\(\sum_{i=l}^{r} A_i = S-K\)

となります。そこで、目標とする区間の合計を \(T=S-K\) と置きます。

まず、\(T\) の値によって次のように場合分けできます。

  • \(T<0\) の場合
    全日程を実施しても費用が \(K\) 円に届かないため、答えは \(0\) です。
  • \(T=0\) の場合
    何も取りやめなければ費用がちょうど \(K\) 円になります。また、\(A_i \ge 1\) なので、空でない区間を取りやめると費用は必ず減ります。したがって、答えは \(1\) です。
  • \(T>0\) の場合
    合計が \(T\) となる空でない連続区間の個数を求めます。

すべての区間を調べると間に合わない

連続区間は \(\frac{N(N+1)}{2}\) 個あります。累積和を使って各区間の合計を \(O(1)\) で求めても、すべて調べるには \(O(N^2)\) 時間が必要です。

\(N\) は最大 \(10^6\) なので、この方法では間に合いません。

各要素が正であることを利用する

\(A_i \ge 1\) なので、区間の合計には次の性質があります。

  • 右端を右に動かして要素を追加すると、合計は必ず増える。
  • 左端を右に動かして要素を取り除くと、合計は必ず減る。

この性質から、右端を順に進めながら、合計が \(T\) を超えたときだけ左端を進めるしゃくとり法が使えます。

また、右端を固定すると、左端を進めるたびに合計は必ず減るため、合計が \(T\) となる左端は高々 \(1\) つです。そのため、右端ごとに一致するかを一度確認すれば十分です。

アルゴリズム

  1. \(T=\sum_{i=1}^{N} A_i-K\) を求めます。
  2. \(T<0\) なら \(0\)、\(T=0\) なら \(1\) を出力して終了します。
  3. \(T>0\) なら、次の変数を用意します。
    • left:現在の区間の左端(\(0\) 始まり)
    • current:現在の区間の合計
    • answer:条件を満たす区間の個数
  4. 配列の要素を左から順に、区間の右端として追加します。
    • 追加した要素を current に加えます。
    • current > T の間、A[left] を引いて left を \(1\) 増やします。
    • current == T なら、answer を \(1\) 増やします。
  5. answer を出力します。

具体例

\(A=[1,2,1,2]\)、\(K=3\) の場合を考えます。

全体の合計は \(S=6\) なので、取りやめる区間の目標合計は \(T=3\) です。

追加する日の番号 追加後の合計 左端を調整した後の区間 合計 答えへの加算
\(1\) \(1\) \([1]\) \(1\) \(0\)
\(2\) \(3\) \([1,2]\) \(3\) \(1\)
\(3\) \(4\) \([2,1]\) \(3\) \(1\)
\(4\) \(5\) \([1,2]\) \(3\) \(1\)

条件を満たす取りやめ方は、日付の区間でいうと \([1,2]\)、\([2,3]\)、\([3,4]\) の \(3\) 通りです。

なぜ数え漏れがないのか

区間の合計が \(T\) を超えているとき、その左端を残したまま右端をさらに進めても、合計は増えるだけです。したがって、その左端は今後も答えにならず、取り除いて構いません。

一方、合計が \(T\) 以下になった時点で、左端をさらに進めると合計はより小さくなります。よって、その右端については、現在の合計が \(T\) と一致するかを確認するだけで十分です。

これをすべての右端について行うことで、条件を満たす区間を重複なく数えられます。

計算量

  • 時間計算量: \(O(N)\)
    右端は \(N\) 回進み、左端も全体で高々 \(N\) 回しか進みません。while 文があっても、処理回数の合計は \(O(N)\) です。
  • 空間計算量: \(O(N)\)
    配列 A を保存するために必要です。しゃくとり法で使う追加の変数は \(O(1)\) です。

実装のポイント

  • \(T=0\) は先に処理します。
    「何も取りやめない」は \(1\) 通りです。空区間を右端ごとに数えてしまわないようにします。

  • 合計が目標を超えている間だけ左端を進めます。
    条件は current > target です。一致した時点の区間を数える必要があるため、>= にしないよう注意してください。

  • 区間が空になる場合もあります。
    追加した要素だけで目標を超える場合、左端を進めて合計が \(0\) になることがあります。\(T>0\) の場合だけこの処理を行うため、空区間が答えに加算されることはありません。

  • 合計の大きさに注意します。
    \(\sum A_i\) は最大 \(10^{15}\) です。Python の整数では扱えますが、固定長整数を使う言語では \(64\) ビット整数が必要です。

  • 入力が大きいため、高速な入力方法を使います。
    提示コードでは sys.stdin.buffer を利用して読み込んでいます。

    ソースコード

import sys


def main():
    input = sys.stdin.buffer
    N, K = map(int, input.readline().split())
    A = list(map(int, input.read().split()))

    target = sum(A) - K
    if target <= 0:
        print(1 if target == 0 else 0)
        return

    left = 0
    current = 0
    answer = 0

    for value in A:
        current += value
        while current > target:
            current -= A[left]
            left += 1
        if current == target:
            answer += 1

    print(answer)


if __name__ == "__main__":
    main()

この解説は gpt-6-astra-medium によって生成されました。