AtCoder159の「D-BannedK」の問題で何が間違っているかわからない

Question

AtCoderのD-BannedKの問題で下記コードを提出したのですが、「WA」の結果となります。
サンプルケースではあっているのですが、一部のテストケースで「WA」となるようです。
何が間違っているのかについて教えてください。
提出結果はこちらです。

""" スニペット """
def get_int():
	return int(input())

def get_ints():
	return list(map(int, input().split()))
""" スニペット """
import math
# 2つ選ぶ通り数を算出 nC2
def choose2(n):
    return math.floor(n*(n-1)/2)

# インプット
N = get_int()
An = get_ints()
""" 全ての要素から2つのボールを選ぶ通り数を求める """
uni = [0] * (N+1)
# ユニーク数の配列を求める
for i in range(N):
    uni[An[i]] += 1
sumWay = 0
# 各数値の2通りの選択通り数を足していく
for i in range(N):
    sumWay += choose2(uni[i])
""" 全ての要素数 - 削除する要素の通り数 + 削除する要素を引いた際の通り数 を求める """
for i in range(N):
    print(sumWay - choose2(uni[An[i]]) + choose2(uni[An[i]]-1)) 

Answer 1

テストケースの中身

AtCoderでは、テストケースの中身が(一部のコンテストを除き)公開されています。
トップページの便利リンク集から飛べます。
今回のテストケースはこれです。

WAになる原因

uni配列では、要素i+1に数字iが書かれたボールの個数が入っています。
そのため、uni[0] = 0です。(なぜなら、1 <= A[i] <= Nなので)
そして、uni[N]には数字Nが書かれたボールの個数が入っています。
しかし、18行目のfor文は0~n-1までしか回っていません。
そこを訂正するとACしました。
少し他にも書き換えた部分があるので、質問の内容とは関係ないですが、説明します。

前計算

何度もchoose2 を計算するのは、(制約が大きくなった場合は)計算量が大きくなってしまいTLEする可能性が高くなってしまいます。
今回は、0C2 ~ nC2を求めたいので、事前にcon配列にそれぞれ代入しておきました。

Comments

[Msk07]

bamboo_music_acさん

ご回答ありがとうございます！
>AtCoderでは、テストケースの中身が(一部のコンテストを除き)公開されています。
テストケースの中身が公開されているなんて知りませんでした。
大変貴重な情報をいただきありがとうございます！
WAになった原因については理解できました。
N個目の要素にアクセスできていなかったのですね。
回答例や、応用問題への方法まで教えていただき大変勉強になりました。
ここまで丁寧でわかりやすいご回答をいただき感動いたしました。
ありがとうございました！！！

[bamboo_music_ac]

競プロは、結構有志の便利なサイトとかコンテンツがいっぱいありますよ〜！

Answer 2

sumWay = 0
# 各数値の2通りの選択通り数を足していく
for i in range(N):
    sumWay += choose2(uni[i])

choose2(uni[N])を足し忘れてます

Comments

[Msk07]

yudedako67さん

ご回答ありがとうございます。
uni[N]の要素が足せていないことを理解できました。
ありがとうございました！