AtCoder Beginner Contest 159 D - Banned K【Python】

https://atcoder.jp/contests/abc159/tasks/abc159_d

AtCoder ProblemsでDifficulty: 496、Solve Probability: 58%でした。

まず全体の書かれている整数が等しいような異なる2つのボールを選び出す方法の数を求め、k番目ごとに全体からAiに関わる組み合わせの数を足し引きして答えを求めました。

from collections import defaultdict

N = int(input())
A = list(map(int, input().split()))

d = defaultdict(int)
for i in range(N):
    d[A[i]] += 1

all = 0
for v in d.values():
    if v >= 2:
        all += v * (v - 1) // 2
ans = 0
for i in range(N):
    ak = A[i]
    cnt = d[ak]
    if cnt == 2:
        print(all - cnt * (cnt - 1) // 2)
    elif cnt > 2:
        print(all - cnt * (cnt - 1) // 2 + (cnt - 1) * (cnt - 2) // 2)
    else:
        print(all)