E - Sandwiches 解説 by climpet

別解

公式解説とは異なり、どちらかというと両端の要素に着目する解法です。

右端の位置 \(k\) を固定したときに、条件を満たす三つ組 \((i, j, k)\) がいくつあるかを考えます。ここで、\(P_k = \{i < k \mid A_i = A_k\}\) (つまり、\(k\) より左にあって値が \(A_k\) に等しい位置の集合) とします。

いま、\(A_i \ne A_j\) という条件を一旦無視してみます。すると、\(i \in P_k\) を一つ取ったとき、\(j\)\(\{i+1, ..., k-1\}\)\(k-1-i\) 通りから自由に選ぶことができます。つまりこのような条件を満たす三つ組 \((i, j, k)\) の個数は、\(\sum_{i \in P_k} (k-1-i) = \sum_{i \in P_k} (k-1) - \sum_{i \in P_k} i = (k-1)|P_k| - \mathrm{sum}(P_k)\) 個となります。

このままだと無視した条件 (\(A_i \ne A_j\)) の分を足しすぎてしまいますが、これは \(A_i = A_j = A_k\) となる三つ組の個数を考えればよく、簡単な計算で \(|P_k|(|P_k| - 1) / 2\) 個であることがわかります。

つまり、\(k\) を固定したとき、条件を満たす三つ組 \((i, j, k)\) の個数は、\((k-1)|P_k| - \mathrm{sum}(P_k) - |P_k|(|P_k| - 1)/2\) 個であることがわかります。この式をよく見ると、集合 \(P_k\) を愚直に持つ必要はなく、要素数 \(|P_k|\) と総和 \(\mathrm{sum}(P_k)\) だけわかれば十分です。

というわけで、\(k\) を左から順に見ていき、各値 \(x\) ごとに出現済みの位置の個数と総和を二つの配列で管理しておけばよいです。計算量は全体で \(O(N)\) 時間となります。

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

#include <iostream>
#include <vector>
using namespace std;

int main(){
  int n;
  cin >> n;
  vector<long long> sum(n + 1), size(n + 1);
  long long ans = 0;
  for (int k = 1; k <= n; ++k) {
    int x;
    cin >> x;
    ans += (k - 1) * size[x] - sum[x] - (size[x] * (size[x] - 1) / 2);
    size[x] += 1;
    sum[x] += k;
  }
  cout << ans << "\n";
}

投稿日時:
最終更新: