あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC159 D - Banned K

競プロ

問題原文

問題要旨

ボールが $N$ 個あり、 $i$ 番目のボールには $A_i$ が書かれている。
$k=1, 2, ... N$ について、以下の問題を解け。

$k$ 番目のボールを除いた $N - 1$ 個のボールから、書かれている整数が同じような異なる２つのボールを選ぶ方法はいくつか？ただし、選ぶ順番は考慮しない。
$1 \leq N \leq 2 × 10^{5}$

解法

各数字が何回出現するかは欲しくなる。
各 $K$ について愚直に各数字が何回出現するかを再計算すると間に合わない。
よく考えると、各 $k$ について変化するのは $A_k$ の出現数のみなので、そこだけ再計算して計算すればいい。

実装

from collections import defaultdict
N = int(input())
A = list(map(int, input().split()))

D = defaultdict(int)
for a in A:
    D[a] += 1

ans_base = 0
for v in D.values():
    ans_base += v * (v - 1) // 2

for a in A:
    print(ans_base - (D[a] * (D[a] - 1) // 2) + ((D[a] - 1) * (D[a] - 2) // 2))

感想

一部だけ再計算が身につきそうな良い典型問題。