【ABC202 C】「Made Up」を解く【Python3】

コンテスト一覧へ

 

長さ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」に足していき、最後にこれを出力します。

【ABC202】解説記事リスト

コンテスト一覧に戻る

コメントを残す

メールアドレスが公開されることはありません。