けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 318 E - Sandwiches (緑色, 450 点)

AtCoder ABC-E 緑色diff 3つのものの真ん中に着目する 3つ組(i<j<k)の問題 累積和 累積和テク:左右両端からの累積和や累積結果を前処理 indexベースで考える 数列 解説あとで詳しく書く 典型要素を詰め合わせた教育的問題 ある量を固定して考える 【問題集】累積和・二分探索法・しゃくとり法 AtCoder450点

色んな解法がありそう。今回は、あまり解説書かずに備忘録程度に。

問題概要

長さ N の数列 A_{1}, A_{2}, \dots, A_{N} がある。各要素は 1 以上 N 以下の整数値である。

  • 1 \le i \lt j \lt k \le N
  • A_{i} = A_{k}
  • A_{i} \neq A_{j}

制約

  • 3 \le N \le 3 \times 10^{5}

メモ

各値ごとに考えていった。つまり、 v = A_{i} = A_{k} として、 v = 1, 2, \dots, N について求めていった。

各 v について、 A_{i} = v となる添字 i を集めてあげて、それをもとに計算していく。全体の計算量は O(N) となる。

コード

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

int main() {
    int N;
    cin >> N;
    vector<int> A(N);
    vector<vector<int>> ind(N+1);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        ind[A[i]].push_back(i);
    }
    
    long long res = 0;
    for (int i = 0; i <= N; ++i) {
        vector<long long> diffs;
        for (int j = 0; j+1 < ind[i].size(); ++j) {
            diffs.push_back(ind[i][j+1] - ind[i][j] - 1);
        }
        long long s = diffs.size();
        
        for (long long j = 0; j < s; ++j) {
            res += diffs[j] * (s - j) * (j + 1);
        }
    }
    cout << res << endl;
}