yukicoder No.118 門松列(2) - masuTomo’s blog

★★に早く慣れたい。

問題

長さが違う3本の竹を選べば適当に並べ替えることで門松列にすることができるので，異なる長さの竹を3本選ぶ選び方の総数（同じ長さの竹も区別する）を求めればよい。
同じ長さの竹が何本あるかがわかればあとは計算できる。
例えば $1,2,3$ という長さの3本の竹を選んだとして，それぞれの竹が $a,b,c$ 本ずつあるとする。このときの門松列の数は $a\times b \times c$ である。
あとは，長さが異なる3本の竹を選ぶ方法はitertoolsのcombinationsを利用すれば簡単である。

itertools --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.8.1 ドキュメント

実装

import itertools
mod = 10**9 + 7
N = int(input())
A = list(map(int,input().split()))
A.sort()
listA = [0 for _ in range(max(A)+1)]
setA = set(A)
for  a in setA:
    listA[a] = A.count(a)
ans = 0
for c in itertools.combinations(setA, 3):
    ans += listA[c[0]]*listA[c[1]]*listA[c[2]]%mod
print(ans%mod)