Official

E - 連続区間の選択 / Selection of Contiguous Intervals Editorial by admin

GPT 5.2 High

Overview

We count the number of intervals \((l,r)\) whose “interval sum + length × M” is at most \(K\). We rephrase this as a “comparison of prefix sums” and efficiently count using a Fenwick tree (BIT).

Analysis

The cost of an interval \((l,r)\) is $\((A_l+\cdots+A_r) + (r-l+1)M\)\( If we consider a new sequence \)B_i = A_i + M\( where \)M\( is added to each element, the cost simply becomes \)\(B_l + B_{l+1} + \cdots + B_r\)\( and the condition becomes equivalent to \)\(B_l+\cdots+B_r \le K\)$

Defining the prefix sums as $\(P[0]=0,\quad P[i]=\sum_{j=1}^{i} B_j = \sum_{j=1}^{i}(A_j+M)\)\( the interval sum is \)\(B_l+\cdots+B_r = P[r]-P[l-1]\)\( so the condition becomes \)\(P[r]-P[l-1] \le K\)\( that is, \)\(P[l-1] \ge P[r]-K\)$

Therefore, the answer is obtained by summing up, for each \(r\), “the number of past prefix sums \(P[0],P[1],...,P[r-1]\) that are at least \(P[r]-K\).”

Why the Naive Approach Fails

Trying all pairs \((l,r)\) is \(O(N^2)\), which for \(N \le 2\times 10^5\) means up to about \(2\times 10^{10}\) pairs — far too slow.

Also, since \(A_i\) can be negative, the interval sum does not necessarily increase monotonically, so the typical “two pointers” technique cannot be used either.

Thus, we need a data structure to count “the number of prefix sum comparisons.”

Algorithm

  1. Compute the prefix sum array \(P[0..N]\) as \(P[i]=P[i-1]+A_i+M\).
  2. We want to use a Fenwick tree to manage “the count of \(P\) values seen so far,” but \(P\) can be as large as \(10^{18}\), so we cannot use the values directly as indices.
    Therefore, we perform coordinate compression: sort the values of \(P\) and remove duplicates to create an array xs.
  3. The Fenwick tree stores “the count (frequency) of prefix sums inserted so far.”
  4. For each \(r=1..N\), do the following:
    • Compute the threshold \(\text{thr} = P[r]-K\).
    • Use bisect_left(xs, thr) to find “how many values in xs are less than \(\text{thr}\) (= the boundary after compression).”
    • Use the Fenwick tree’s prefix sum to find “the number of past \(P\) values less than \(\text{thr}\)” (cnt_lt).
    • If the total number of values inserted so far is total, then the number satisfying \(P[l-1]\ge \text{thr}\) is $\(\text{total} - \text{cnt\_lt}\)$ so add this to the answer.
    • Finally, add \(P[r]\) to the Fenwick tree and increment total.
  5. As an initial state, insert \(P[0]=0\) beforehand (to handle the case \(l=1\)).

With this method, we can find “the number of valid \(l\) values for each \(r\)” in \(O(\log N)\).

Complexity

  • Time complexity: \(O(N\log N)\) (sorting for coordinate compression \(O(N\log N)\) + Fenwick tree operations for each \(r\) \(O(\log N)\))
  • Space complexity: \(O(N)\) (prefix sum array, compressed array, Fenwick tree)

Implementation Notes

  • \(P[i]\) and \(K\) can be as large as \(10^{18}\), so use 64-bit integers (in Python, the standard int is fine).

  • The Fenwick tree is 1-indexed, so add +1 to the compressed index before calling add.

  • Instead of directly counting “\(P[l-1] \ge P[r]-K\),”
    it is easier to handle with the Fenwick tree’s prefix sum by subtracting “the count of \(P[l-1] < P[r]-K\).”

    Source Code

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()

This editorial was generated by gpt-5.2-high.

posted:
last update: