E - 連続区間の選択 / Selection of Contiguous Intervals Editorial by admin
GPT 5.2 High概要
各区間 \((l,r)\) の「区間和+長さ×M」が \(K\) 以下になるものを数えます。これを「累積和の大小比較」に言い換え、Fenwick木(BIT)で高速に数え上げます。
考察
区間 \((l,r)\) のコストは
$\((A_l+\cdots+A_r) + (r-l+1)M\)\(
です。ここで各要素に \)M\( を足した新しい列 \)B_i = A_i + M\( を考えると、コストは単に
\)\(B_l + B_{l+1} + \cdots + B_r\)\(
になり、条件は
\)\(B_l+\cdots+B_r \le K\)$
と同値になります。
累積和を $\(P[0]=0,\quad P[i]=\sum_{j=1}^{i} B_j = \sum_{j=1}^{i}(A_j+M)\)\( とすると、区間和は \)\(B_l+\cdots+B_r = P[r]-P[l-1]\)\( なので条件は \)\(P[r]-P[l-1] \le K\)\( すなわち \)\(P[l-1] \ge P[r]-K\)$ となります。
よって「各 \(r\) について、過去の累積和 \(P[0],P[1],...,P[r-1]\) のうち \(P[r]-K\) 以上の個数」を足し上げれば答えです。
素朴解がダメな理由
全ての \((l,r)\) を試すと \(O(N^2)\) で、\(N \le 2\times 10^5\) では最大約 \(2\times 10^{10}\) 通りになり到底間に合いません。
また、\(A_i\) は負になり得るので、区間和が単調に増えるとは限らず、典型的な「尺取り法(two pointers)」も使えません。
そこで「累積和の大小関係の個数」をデータ構造で数える必要があります。
アルゴリズム
- \(P[i]=P[i-1]+A_i+M\) として累積和配列 \(P[0..N]\) を作る。
- これから Fenwick 木で「これまでに出てきた \(P\) の値の個数」を管理したいが、\(P\) は最大で \(10^{18}\) 規模になり得るため、そのまま添字にできません。
そこで 座標圧縮を行い、\(P\) の値をソートして重複を除いた配列xsを作る。 - Fenwick木には「これまでに挿入した累積和の個数(頻度)」を入れる。
- 各 \(r=1..N\) について以下を行う:
- 閾値 \(\text{thr} = P[r]-K\) を計算。
bisect_left(xs, thr)により「\(\text{thr}\) 未満の値がxs内で何個あるか(=圧縮後の境界)」を得る。- Fenwick木の prefix sum を使って「過去の \(P\) のうち \(\text{thr}\) 未満の個数」を求める(
cnt_lt)。 - これまでの総数を
totalとすると、条件 \(P[l-1]\ge \text{thr}\) を満たす個数は
$\(\text{total} - \text{cnt\_lt}\)$
なのでそれを答えに加算する。 - 最後に \(P[r]\) を Fenwick木に追加し、
totalも増やす。
- 初期状態として \(P[0]=0\) を先に挿入しておく(\(l=1\) を扱えるようにするため)。
この方法で「各 \(r\) に対して条件を満たす \(l\) の数」を \(O(\log N)\) で求められます。
計算量
- 時間計算量: \(O(N\log N)\)(座標圧縮のソート \(O(N\log N)\) + 各 \(r\) のFenwick操作 \(O(\log N)\))
- 空間計算量: \(O(N)\)(累積和・圧縮配列・Fenwick木)
実装のポイント
\(P[i]\) や \(K\) は \(10^{18}\) まで行き得るので、64bit整数で扱う(Pythonなら通常の
intでOK)。Fenwick木は 1-indexed なので、圧縮した index に
+1してからaddする。「\(P[l-1] \ge P[r]-K\)」の個数を直接数える代わりに、
「\(P[l-1] < P[r]-K\) の個数」を引く形にすると Fenwick の prefix sum で処理しやすい。ソースコード
import sys
from bisect import bisect_left
class Fenwick:
def __init__(self, n):
self.n = n
self.bit = [0] * (n + 1)
def add(self, i, v):
while i <= self.n:
self.bit[i] += v
i += i & -i
def sum(self, i):
s = 0
while i > 0:
s += self.bit[i]
i -= i & -i
return s
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
N, M, K = data[0], data[1], data[2]
A = data[3:]
P = [0] * (N + 1)
for i in range(1, N + 1):
P[i] = P[i - 1] + A[i - 1] + M
xs = sorted(set(P))
fw = Fenwick(len(xs))
ans = 0
total = 0
# insert P[0]
fw.add(bisect_left(xs, P[0]) + 1, 1)
total += 1
for r in range(1, N + 1):
thr = P[r] - K
idx = bisect_left(xs, thr) # number of values < thr
cnt_lt = fw.sum(idx)
ans += total - cnt_lt
fw.add(bisect_left(xs, P[r]) + 1, 1)
total += 1
print(ans)
if __name__ == "__main__":
main()
この解説は gpt-5.2-high によって生成されました。
posted:
last update: