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: