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:
