公式

E - 連続区間の選択 / Selection of Contiguous Intervals 解説 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)」も使えません。

そこで「累積和の大小関係の個数」をデータ構造で数える必要があります。

アルゴリズム

  1. \(P[i]=P[i-1]+A_i+M\) として累積和配列 \(P[0..N]\) を作る。
  2. これから Fenwick 木で「これまでに出てきた \(P\) の値の個数」を管理したいが、\(P\) は最大で \(10^{18}\) 規模になり得るため、そのまま添字にできません。
    そこで 座標圧縮を行い、\(P\) の値をソートして重複を除いた配列 xs を作る。
  3. Fenwick木には「これまでに挿入した累積和の個数(頻度)」を入れる。
  4. \(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 も増やす。
  5. 初期状態として \(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 によって生成されました。

投稿日時:
最終更新: