PythonでABC159-Dを解く

はじめに

Boot camp for Beginnersを解きます。

ABC159-D Banned K

考えたこと
$O(N)$くらいに抑えないと通らないので、それぞれの$k$に対してコンビネーションを計算していては間に合わない。
少し考えてみると、種類の違う数字の組合せは独立であることが分かる。例えば、取り除かれたボールが1のとき、1以外の組合せの総和は変化しない。では、取り除かれた数字の組み合わせはどう変化するのかを考える。$n$個から2個選ぶ組合せは$\frac{n(n-1)}{2}$で、$n-1$から2個選ぶ組合せは$\frac{(n-1)(n-2)}{2}$になるから差を取ると$n-1$となる。
あとは、Counterで個数を調べて全体の組み合わせを調べて計算する。

from collections import Counter
n = int(input())
a = list(map(int,input().split()))

c = Counter(a)
key = c.keys()
comb = 0
for i in key:
    comb += (c[i]) * (c[i]-1) // 2 #禁止されていない状態の組み合わせを調べる

for i in a:
    ans = comb - (c[i]-1) #a[i]が禁止されている状態の組み合わせ
    print(ans)

まとめ

2ヶ月前の自分が解けなかった問題が解けるようになっていて成長を感じました。ではまた。
余談ですが、今日はあーだこーだーがありますよー。