AtCoder 勉強記録 ~茶色コーダーを目指して~ [ABC 159-D]

今回の問題

AtCoder Beginner Contest 159 D - Banned K
atcoder.jp

私の解答

https://atcoder.jp/contests/abc159/submissions/12682578

import collections
n = int(input())
a = list(map(int,input().split()))
new_a = collections.Counter(a) // 出現回数の分布取得
 
total = 0
for key, val in new_a.items():
    if val > 1:
        total += (val * (val - 1) // 2) // すべてのボールがある状態の取り出す方法の数
 
 
for i in a:
    ans = total
    ans -= new_a[i]-1 // 出現回数を1つ減らした取り出し方法数をtotalから引く
    print(ans)

解説

普通にループで毎回取り出し方法の数を数えるとTLEになってしまう。
考え方としては、最初にすべてのボールがある状態の取り出す方法の数を計算しておいて、
そこからそれぞれ出現回数を1つ減らした取り出し方法数を引けばよい。

計算量の想定ができないから無駄な実装をしてしまいました。もう少し慣れが必要そう。
ちなみに下記がTLEになった実装です。

import collections
import itertools
n = int(input())
a = list(map(int,input().split()))
 
for i in range(n):
    ans = 0
    pop = a.pop(i)
    new_a = collections.Counter(a)
    # print(new_a)
    for key,val in new_a.items():
        # print(val)
        if val > 1:
            ans += (val * (val - 1) // 2)
    a.insert(i, pop)
    print(ans)