2014 Yandex.Algorithm Elimination Stage, Round 3 C - Intervals

解法

しゃくとり法で解いた.
条件を満たさないようなペアの数を数えて,全体から引くことで求める.
しゃくとりの過程で,今現在の区間が overlap しうるので,その分は足すことで重複して引かないようにしている.

条件を満たさないような任意のペア (l, r) が,必ずしゃくとりの過程でカウントされることを証明する必要がある.
カウントされないような動きとしては,今見ている区間の右端が l と r の間にあるのに,左端を狭めるときに l と r の間に入り込んでしまう場合であり,これしかない.
しかし,右端が l と r の間にあり,かつ左端を狭めていく時は,必ず l か l の前で停止する.(l, r) の間に条件を満たすペアが無いことからこれは明らかである.
よって,任意の条件を満たさないペア (l, r) はしゃくとりの過程で必ずカウントされる.

ソースコード

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

using ll = long long;

int main() {
    int n, d;
    scanf("%d %d", &n, &d);
    vector<int> a(n);
    for(int i = 0; i < n; ++i) {
        scanf("%d", &a[i]);
    }

    unordered_map<int, int> cnt;
    int l = 0, r = 0;
    int prev_r = -1;
    ll ans = 1LL * n * (n - 1) / 2;
    while(r < n) {
        while(r < n && cnt[a[r] + d] == 0 && cnt[a[r] - d] == 0) {
            cnt[a[r++]] += 1;
        }
        const ll width = r - l;
        const ll overlap = max(0, prev_r - l);
        ans -= width * (width - 1) / 2;
        ans += overlap * (overlap - 1) / 2;
        while(cnt[a[r] + d] > 0 || cnt[a[r] - d] > 0) {
            cnt[a[l++]] -= 1;
        }
        prev_r = r;
    }

    cout << ans << endl;
}

感想

この解法で解けそうだな~と思って通してしまった….
証明はあとづけである.反省.しゃくとり書く時いつもこんな感じだ….