AtCoder茶色解く２

今回解くのはこちら
D - Banned K

Nの数は2*10^5が上限
Aの数は1~Nまで

まずは簡単にするために、すべてのAを数え上げて記録する

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

for a in A:
  if a not in di.keys(): # diに何もなかった場合
    #print("add")
    di[a] = 1
  else:
    #print("sum")
    di[a] += 1
#print(di)

さらに、得られたまとめデータをもとに各Aに対してfor文を回す

# 組み合わせ計算
sum = 0
for k in A:
  for a, b in di.items():
    if a == k: b -= 1
    #print(a,b)
    sum += int((b*(b-1))/2*1)
  print(sum)
  sum = 0

一部がTLE出ている
よって、最適化を行っていく

まずは、計算結果をメモするように変更した

# 組み合わせ計算
sum = 0
ll = {}
for k in A:
  if k not in ll.keys():
    for a, b in di.items():
      if a == k: b -= 1
      #print(a,b)
      sum += int((b*(b-1))/2*1)
    print(sum)
    ll[k] = sum
  else:
    print(ll[k])
  sum = 0

これで一部高速化でき、ACの数は増えたが、以前TLEが存在する状況であった。

ここで、以下記事を参考に考え方を変えた
AtCoder ABC 159 D - Banned K (400 点) - けんちょんの競プロ精進記録

愚直に毎回の数列について計算したのでは O(N2) の計算量を要してしまう。今回のような、各 k について、それを除いたものについて考えるタイプの問題では

まず全体の数列についての答えを求めておいて
各要素を取り除くことによる「影響」を計算し、
その影響分を引く

先に組み合わせ数の計算値を求めておき、そこから各パターンを追加で計算する。
差分更新というやり方らしい

これで書き換えたものが以下

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

for a in A:
  if a not in di.keys(): # diに何もなかった場合
    #print("add")
    di[a] = 1
  else:
    #print("sum")
    di[a] += 1

# 全体量計算
sum = 0
for a, b in di.items():
  #print(a,b)
  sum += int((b*(b-1))/2*1)
  
# 組み合わせ計算
for k in A:
  n = di[k]
  bef = n*(n-1)/2
  aft = (n-1)*(n-2)/2
  print(int(sum-bef+aft))

無事ACが取れました。