エイシングプログラミングコンテスト2021（ABC202）を振り返る C問題

C - Made Up

C - Made Up

問題

1以上N以下の整数からなる長さNの数列A,B,Cが与えられる、1以上N以下の整数i,jの組(i,j)であって、Ai=Bcjとなるものの総数を求めてください

入力例1

3
1 2 2
3 1 2
2 3 2

Aの数列から考えるパターンと、Cの数列から考えるパターンがあるようなんですが、私はCの数列から考えました。

Cの数列の１番目は2で、Bの数列の2番目1はAに1つ含まれています。
Cの数列の２番目は3で、Bの数列の3番目2はAに2つ含まれています。
Cの数列の３番目は2で、Bの数列の2番目1はAに1つ含まれています。
ということで、上記入力例の場合の出力は４となります。

TLEになった回答

上記をそのままコードにしたのが以下です。

n = int(input())
a = [int(i) for i in input().split()]
b = [int(i) for i in input().split()]
c = [int(i) for i in input().split()]
ans = 0
for i in c:
   ans += a.count(b[i-1])
print(ans)

TLE（Time Limit Exceeded）が出まくりで駄目でした。思うにループの中の「count」メソッドが都度ループ回してるんじゃないかと思うんですよね。

ACできた回答

ということで、事前に数列に含まれる要素とその個数を確認するようにしたのが以下です。

import collections

n = int(input())
a = [int(i) for i in input().split()]
b = [int(i) for i in input().split()]
c = [int(i) for i in input().split()]
a_count = collections.Counter(a)
c_count = collections.Counter(c) 
ans = 0
for key, val in c_count.items():
   i = b[key-1]
   if i in a_count:
       ans += a_count[i] * val
print(ans)

「collections.Counter」と言う標準ライブラリを使ってますが、他の人の回答見たら数列に含まれる数字は1以上N以下であることが保証されているので、

a_count = [0]*n
for i in a:
   a_count[i-1] += 1

とかやってリストだけで実装することも可能でしたね。

最初は時間超過とかで正解できなかった問題をコードを書き換えることで正解に持っていくっていうのが、面白いです（正解した場合のみ）