AtCoder Beginner Contest 159 [ D - Banned K ] https://atcoder.jp/contests/abc159/tasks/abc159_d

abc159_d.py

'''
[問題]
https://atcoder.jp/contests/abc159/tasks/abc159_d

[参考]
https://blog.hamayanhamayan.com/entry/2020/03/22/235319
　C(個数,2)、つまり、個数×(個数 - 1)/2の総和を取れば答えになる。

[結果]
PyPy3(7.3.0) AC 249ms
Python(3.8.2) AC 290ms
'''

import sys
from collections import defaultdict
import math

'''
こっちは、真面目に計算してて、遅い！
def combinations_count(n, r):
    return math.factorial(n) // (math.factorial(n - r) * math.factorial(r))
'''

# r = 2固定なので、決め打ちでこの計算でOK
def combinations_count(n):
    return n*(n-1)//2

sys.setrecursionlimit(10 ** 6)  # 再帰上限の引き上げ
input = sys.stdin.readline

INF = 2 ** 63 - 1

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

d = defaultdict(int)

for a in A:
    d[a] += 1

a_tmp = defaultdict(int)
a_sum = 0
for k, v in d.items():
    if v < 2:
        continue
    # 組み合わせ計算が遅くてTLE
    #a_tmp[k] = combinations_count(v, 2)

    # こっちの最適化した計算でやればAC
    a_tmp[k] = combinations_count(v)
    a_sum += a_tmp[k]


for i in range(N):

    temp = d[A[i]] - 1
    if temp < 2:
        print(a_sum - a_tmp[A[i]])
    else:
        #sa = a_tmp[A[i]] - combinations_count(temp, 2)

        # こっちの最適化した計算でやればAC
        sa = a_tmp[A[i]] - combinations_count(temp)
        print(a_sum - sa)