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\) つです。そのため、右端ごとに一致するかを一度確認すれば十分です。
アルゴリズム
- \(T=\sum_{i=1}^{N} A_i-K\) を求めます。
- \(T<0\) なら \(0\)、\(T=0\) なら \(1\) を出力して終了します。
- \(T>0\) なら、次の変数を用意します。
left:現在の区間の左端(\(0\) 始まり)current:現在の区間の合計answer:条件を満たす区間の個数
- 配列の要素を左から順に、区間の右端として追加します。
- 追加した要素を
currentに加えます。 current > Tの間、A[left]を引いてleftを \(1\) 増やします。current == Tなら、answerを \(1\) 増やします。
- 追加した要素を
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 によって生成されました。