11591번: VUDU

https://www.acmicpc.net/problem/11591


$O(n\lg n)$

s[i]: a[1]+a[2]+...+a[i]

(s[i]-s[j])/(i-j)>=p (j<i) 에서
s[j]-pj<=s[i]-pi

매 i마다 위 조건을 만족하는 j는 펜윅 트리로 구할 수 있다.

#include<cstdio>
#include<algorithm>
using namespace std;
typedef long long ll;
const int MXN = 1e6;
int n, sz, bit[MXN + 2];
ll s[MXN + 1], idx[MXN + 1], p, r;
int main() {
        scanf("%d", &n);
        for (int i = 1; i <= n; i++) scanf("%lld", s + i), s[i] += s[i - 1];
        scanf("%lld", &p);
        for (int i = 0; i <= n; i++) idx[i] = s[i] - i*p;
        sort(idx, idx + 1 + n);
        sz = unique(idx, idx + 1 + n) - idx;
        for (int i = 0; i <= n; i++) {
                int t = lower_bound(idx, idx + sz, s[i] - i*p) - idx + 1;
                for (int h = t; h; h -= h&-h) r += bit[h];
                for (int h = t; h <= sz; h += h&-h) bit[h]++;
        }
        printf("%lld", r);
        return 0;
}