AtCoder Beginner Contest 206 C - Swappable - おじさんの競プロ記録

問題

N個の整数からなる数列 $A= \left( A_1, A_2, ..., A_N \right)$ が与えられるので、次の条件を全て満たす整数組 $\left( i, j \right)$ の数を求めてください。

$1 \leq i \lt j \leq N$
$A_i \neq A_j$

制約

入力は全て整数
$2 \leq N \leq 3 \times 10^5$
$1 \leq A_i \leq 10^9$

考えたこと

制約から、全ての整数組を試すのでは間に合いません。ですが、一度それで考えてみて改善できるか確認してみます。

入力例２に注目しました。10個の違う数が含まれた数列では答えは45です。 $A_1$ に対して残りすべてが条件に当てはまるので9個。 $A_2$ は8個と続きます。数列に含まれる数が全て異なるならば、Nとループインデックスで計算できそうです。

数列に等しい数が含まれる場合を考えます。 $A_i$ に注目した時に $A_i = A_j$ になる数がいくつあるかわかればよさそうです。事前に、数列に数が出現する回数を数えておきます。mapを使いました。

最終的には数列の先頭から見ていきます。数列に含まれる数が全て異なる場合から、事前に調べておいた等しい個数を引いてから答えに足していきました。

二重ループから、2回のループになり制限時間に間に合いました。

コード

#define _GIBCXX_DEBUG
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(v) v.begin(), v.end()
using namespace std;
using ll = long long;


int main(void) {
    ll N;
    cin >> N;
    vector<ll> A(N);
    rep(i, N) cin >> A[i];

    map<ll, int> mp;
    rep(i, N) mp[A[i]]++;

    ll ans = 0;
    for (ll i = 0; i < N - 1; ++i) {
        ll residue = N - 1 - i;
        mp[A[i]]--;
        ans += residue - mp[A[i]];
    }

    cout << ans << endl;
    
    return 0;
}

解説を読んで

整数組 $\left( i, j \right)$ が $\frac{N \times \left(N - 1 \right)}{2}$ が、パッと浮かばないのがつらい。