問題概要

$N$ 個の整数 $A_{1}, \dots, A_{N}$ が与えられる。以下の条件を満たす整数 $(i, j)$ の組の個数を求めよ。

$1 \le i \lt j \le N$
$A_{i} \neq A_{j}$

制約

$2 \le N \le 3 \times 10^{5}$
$1 \le A_{i} \le 10^{9}$

考えたこと

まず、 $A_{i} \neq A_{j}$ という条件がないバージョンを考えてみよう。そのときは単に $N$ 個のものから 2 個選ぶ場合の数を求めよ、という問題になる。

${}_{N}{\rm C}_{2} = \frac{N(N-1)}{2}$ 通り

となる。これは AtCoder では定期的によく出されている。

drken1215.hatenablog.com

補集合を考える

さていよいよ、 $A_{i} \neq A_{j}$ という条件を考えていこう。しかしこの条件を考えるのは難しいので、かわりに「逆」の方を数えることにしよう。つまり、逆に

$1 \le i \lt j \le N$
$A_{i} = A_{j}$

であるような場合の数を数えることにしよう。その答えを $P$ としたとき、 $\frac{N(N-1)}{2} - P$ を出力すれば OK。

$A_{i}$ の値ごとに分類

さて、逆の方は数えやすい。たとえば $A = (1, 1, 2, 3, 3, 1, 1)$ だと

1 が 4 個
2 が 1 個
3 が 2 個

となる。 $A_{i} = A_{j} = 1$ となるパターンは「4 個から 2 個を選ぶ場合の数」なので ${}_{4}{\rm C}_{2} = 6$ 通りで、 $A_{i} = A_{j} = 2$ となるパターンは ${}_{1}{\rm C}_{2} = 0$ 通りで、 $A_{i} = A_{j} = 3$ となるパターンは ${}_{2}{\rm C}_{2} = 1$ 通りとなる。これらを合計して $6 + 0 + 1 = 7$ 通りとなる。

一般に、 $A_{i}$ の値を数値ごとに分類して、その個数を $a$ として ${}_{a}{\rm C}_{2}$ の総和を取っていけば OK。

コード

計算量は $A$ の値を分類するときに map などを用いた場合、 $O(N \log N)$ となる。

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

int main() {
    long long N;
    cin >> N;
    map<long long, long long> ma;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        ma[A[i]]++;
    }

    long long P = 0;
    for (auto it: ma) {
        long long num = it.second;
        P += num * (num - 1) / 2;
    }
    cout << N * (N - 1) / 2 - P << endl;
}