ABC369 C問題〜ランレングス圧縮〜

解説動画

動画内の説明で利用していたファイル

コード

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i, n) for (int i = 0; i < (n); ++i) int main() { int n; cin >> n;
    vector a(n);
    rep(i, n) cin >> a[i];

    if (n == 1) {
        cout << 1 << endl;
        return 0;
    }

    vector b(n - 1);
    rep(i, n - 1) {
        b[i] = a[i + 1] - a[i];
    }

    vector<pair<int, int>> vec(1, {b[0], 1});

    for (int i = 1; i < n - 1; i++) {
        if (b[i - 1] == b[i]) {
            vec.back().second++;
        } else {
            vec.push_back({b[i], 1});
        }
    }

    // 確認用のcout
    // for (auto e : vec) {
    //     cout << e.first << " " << e.second << endl;
    // }

    ll ans = n;

    for (auto e : vec) {
        ll count = e.second;
        ans += ((count + 1) * count) / 2;
    }

    cout << ans << endl;

    return 0;
}