[AtCoder] ABC 159 D – Banned K

2020年3月23日2020年12月15日

問題

方針

\( k \) 番目のボールを除いたときに整数が等しいような異なる \( 2 \) つのボールを選び出す方法の数が変化する整数は、\( A_k \) だけなので、全体の組み合わせから \( A_k \) の値を引くような方針で考えます。\( b(x) \) を整数 \( x \) の個数とすると、整数 \( x \) の選び方は、

\[ \dfrac{b(x)(b(x) – 1)}{2}\]

となります。したがって、\( N \) 個のボールに対する整数が等しいような異なる \( 2 \) つのボールを選び出す方法の数を \( s \) とすると、\( k \) 番目のボールを取り除くと、

\[s – \dfrac{b(A_k)(b(A_k) – 1)}{2} + \dfrac{(b(A_k) – 1)(b(A_k) – 2)}{2}\]

となります。

コード

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

const int l = 2000001;
ll b[l]{};

ll func(ll n) {
    return n * (n - 1) / 2;
}

int main() {
    int N;
    cin >> N;
    int A[N];
    for (int i = 0; i < N; i++) {
        cin >> A[i];
        b[A[i]]++;
    }
    ll sum = 0;
    for (int i = 1; i < l; i++) {
        sum += b[i] * (b[i] - 1) / 2;
    }
    for (int i = 0; i < N; i++) {
        cout << sum - func(b[A[i]]) + func(b[A[i]] - 1) << "\n";
    }
    return 0;
}