ABC206 C問題 Swappable をNのオーダーで解く - 競技プログラミング悪戦苦闘日記

ABC206のC問題の解いた過程を書きながら解説します。
公式解説がかなり丁寧なので、厳密な解説はそちらを参考にしてください。

問題の概要

N個の整数からなる配列から、2つの異なる値の組み合わせが何個あるか求める問題です。

難しいポイント

$N \leq 3 \times 10^5$ なので、愚直にループを回すとタイムオーバーします。
同じ数が配列に含まれるときに重複をどうやって取り除くかがポイントです。

方針

まず全体から2つ選ぶ組み合わせは ${}_N C_2 = \frac{N \times (N - 1)}{2}$ です。
重複を削除する方針を考えるために、愚直に具体例をイメージして書き出してみます。
ここでは、 $[1, 7, 1, 1, 2, 2]$ という配列で表を描いてみます。

1は3個あります。1が重複することによって、黄色の3パターンが条件を満たさないパターンです。
2は2個あります。2が重複することによって、水色の1パターンが条件を満たさないパターンです。
この表を見ると、重複している数字の個数を $v$ としたときに、条件を満たさないパターンの数は ${}_v C_2 = \frac{v \times (v - 1)}{2}$ ということがわかります。
もっと大きい重複数で書き出すともっとイメージがつかめるかもしれません。

実装

collections.Counterを使って、dictionaryに各数字の出現個数を記録しました。
各数字の出現個数から条件を満たさない個数を算出して足していきます。
全体の場合の数から条件を満たさない個数を引いた数を出力します。

from collections import Counter
N = int(input())
a_arr = list(map(int, input().split()))

# 各数字の出現回数を記録する
rui_dict = Counter(a_arr)
minus = 0
for v in rui_dict.values():
    # 重複数を数える
    minus += (v * (v - 1)) // 2
# 全パターンから条件を満たさないパターン数を引いた数が答え
ans = N * (N - 1) // 2 - minus
print(ans)

こんな簡単な問題文なのに、思い込みで問題を読み違えていて時間ロスしました。。。