Official

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


\(\begin{aligned} f(l,r) &= \sum_{i=l}^{r}A_i+(r-l+1)\times M \\ &= \sum_{i=l}^{r}(A_i+M) \end{aligned}\)

となるので、 \(B_i = A_i+M\) として \(A\) の代わりに \(B\) を考えることで、「和が \(K\) 以下になる区間は何個?」という問題になります。

さらに累積和 \(S_n = \sum_{i=1}^{n}B_i\) を考えることで問題は「 \(0 \leq l<r \leq N\) を満たす組 \((l,r)\) のうち、\(S_r-S_l\leq K\) となるものは何個?」となります。

\(r\) を固定した問題を考えると「\(S_0, \ldots,S_{r-1}\) のうち、\(S_r-K\) 以上のものは何個?」 となります。よって \(r\) の昇順に考えることで、元の問題に答えるためには次の操作を高速に行えるデータ構造があれば十分であることがわかります:

  • 要素を \(1\) つ追加する
  • 追加済みの要素のうち \(x\) 以上のものが何個あるかを求める

これは \(S_i\) を予め座標圧縮したうえで、セグメントツリーを用いることで高速に処理することができます。

計算量は \(O(N\log N)\) です。

実装例 (C++)

#include<bits/stdc++.h>
#include<atcoder/fenwicktree>
using namespace std;

int main(){
  int n, m;
  long long k;
  cin >> n >> m >> k;
  vector<int>a(n);
  for(int i=0; i<n; i++) cin >> a[i];
  vector<long long>s(n+1);
  for(int i=0; i<n; i++) s[i+1] = s[i] + a[i] + m;

  vector<long long>d(s.begin(), s.end());
  sort(d.begin(), d.end());
  d.erase(unique(d.begin(), d.end()), d.end());
  auto get_index = [&](long long x){
    return lower_bound(d.begin(), d.end(), x) - d.begin();
  };

  atcoder::fenwick_tree<int>seg(d.size());
  long long ans = 0;
  for(int i=0; i<n+1; i++){
    ans += i - seg.sum(0, get_index(s[i]-k));
    seg.add(get_index(s[i]), 1);
  }
  cout << ans << endl;
}

この他、python では SortedSet もこの操作を高速に行うことができるデータ構造となっていますが、Set の名の通り重複する要素を持つことができません。そこで以下の実装例では、値が重複する要素を複数持てるよう、値と index のペアを管理する実装となっています。

実装例 (Python)

from sortedcontainers import SortedSet

N, M, K = map(int, input().split())
A = list(map(int, input().split()))
S = [0]
for a in A:
  S.append(S[-1] + a + M)

T = SortedSet()
ans = 0
for i, s in enumerate(S):
  ans += i - T.bisect_left((s-K, 0))
  T.add((s, i))
print(ans)

posted:
last update: