長さNの数列A、数列B、数列Cに対して、整数組(i, j)であり、Ai=BCj となるものの数を求める問題です。
提出
import collections
n = int(input())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
c = list(map(int, input().split()))
bb = []
for i in c:
bb.append(b[i-1])
aco = collections.Counter(a)
bco = collections.Counter(bb)
ans=0
for i in range(1,n+1):
ans += aco[i]*bco[i]
print(ans)
collectionsモジュールを使って解いていきます。
まず、a, b, c の入力用のリストを作り、それとは別に、bb の空リストを作成します。
bbの中に、BCiとなる整数を入れておき、aとbb の中で要素が等しいものの数を数えていきます。
collections の Counter()を使うと、それぞれの要素の数がカウントされます。
上記で求めた「aco」「bco」は以下のようになります。
aco = collections.Counter(a)
bco = collections.Counter(bb)
print(aco)
# > Counter({2: 2, 1: 1})
print(bco)
# > Counter({1: 2, 2: 1})
Counter() は、辞書型のサブクラスで、辞書型と同じ操作ができるため、aco[i] の i にキーの数を入れると、個数を出力します。
例えば、上記の bco[2] は 1 になります。
制限として、数列内の数字はどれも、1以上N以下の整数であるため、同じくキーも1以上N以下となります。
組み合わせの数は、それぞれのキーに対して、aco[i]×bco[i] で求められます。
この i を1~Nの範囲で求めて、「ans」に足していき、最後にこれを出力します。