atcoder ABC 318 E - Sandwichesの勉強

atcoder.jp


・参考
THIRD プログラミングコンテスト 2023 アルゴ(AtCoder Beginner Contest 318) - YouTube

・説明 
 参考動画をPythonで実装した。真ん中だけ違うもの(動画のX?X) - 全部同じもの(動画のXXX)を求める。
 また、参考動画での \begin{eqnarray}\sum_{i=1}^{n}(k-i-1)&=&(k-1-1)+(k-2-1)+..+(k-n-1)&=&\\k*n-n-\sum_{i=1}^{n} i&=&(k-1)*n-\sum_{i=1}^{n} i\end{eqnarray}となる。ここでのnはiの個数である。本文のnではない。

from collections import Counter
n=int(input())
a=list(map(int,input().split()))

cnt=[0 for _ in range(n+1)]   #iの個数(1,2,3..,nが何個か).
sum_idx=[0 for _ in range(n+1)]   #iの合計(1のindexの和、1のindexの和、、、nのindexの和).

#mC3の組み合わせ.
def mc3(m):
    return (m*(m-1)*(m-2))//6

#全部同じもの個数を求める。
c=Counter(a)
all_same=0     #全部同じもの個数
for val in c.values():
    all_same+=mc3(val)
    #print(c,val,all_same)


#真ん中だけ違うもの.
tmp=0   #真ん中だけ違うものの個数.
for k in range(n):
    tmp+=(k-1)*cnt[a[k]]-sum_idx[a[k]]
    cnt[a[k]]+=1  
    sum_idx[a[k]]+=k
print(tmp-all_same)