E - Sandwiches Editorial by carrot46

別解

\(x\) を固定し、\(A_i = A_k = x\) かつ \(A_j \neq x\) となる \((i, j, k)\) の個数を数えます。値が \(x\) である要素の個数を \(n_x\) とし、それらの要素の添字を順に \(p_1, \dots, p_{n_x}\) とします。ここで、\(p_m < j < p_{m + 1}\) を満たすような \(j\) について考えると、\(i\)\(p_1, \dots, p_m\)\(m\) 通り、\(k\)\(p_{m + 1}, \dots, p_{n_x}\)\(n_x - m\) 通りの選び方があります。よって、このような \(j\) を固定した時の \((i, j, k)\) の選び方は \(m(n_x - m)\) 通りです。

ゆえに、\(A_i = A_k = x\) かつ \(A_j \neq x\) となる \((i, j, k)\) の個数は \(\sum_{m = 1}^{n_x - 1} (p_{m + 1} - p_m - 1) m(n_x - m)\) により計算出来ます。この式を基に各 \(x\) に対して愚直に計算しても、合計の計算量は \(O(\sum_{x = 1}^N n_x) = O(N)\) となり、十分高速に解くことが出来ます。

実装例 : https://atcoder.jp/contests/abc318/submissions/45200908

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

int main() {
    int N;
    cin>>N;
    vector<int> a(N);
    vector<vector<ll>> pos(N);
    for(int i = 0; i < N; ++i){
        cin>>a[i];
        --a[i];
        pos[a[i]].push_back(i);
    }
    ll res = 0;
    for(auto p : pos){
        ll n = size(p);
        for(int m = 0; m < n - 1; ++m){
            res += (p[m + 1] - p[m] - 1) * (m + 1) * (n - (m + 1));
        }
    }
    cout<<res<<endl;
}

posted:
last update: