Atcoder Beginner Contest 202 C問題の計算量削減させたい(Atcoder,Python3)

Question

最近競技プログラミングを始めてB問題は安定して解けるようになったものです。

最近の課題としてC問題以降の問題での計算量削減があり、ただ普通に解いても桁が大きくなると時間内に処理が終わらずにエラーが起きてしまいます。

今回も解くことができましたがどうしても計算量をどのように減らすかわからずに完答することができませんでした。

今回書いたコードを載せますのでどのような部分を改善したら時間内（2 sec）に処理が終わるのかご教示お願い致します。

個人的にはやはり二重ループや全探索による探索が大きな負担になっていると思うのですが今回はこれ以外での実装が困難だったため二重ループを利用しました。

https://atcoder.jp/contests/abc202/tasks/abc202_c

s = int(input())
lst = []
cnt = 0
for _ in range(3):
	X = list(map(int, input().split()))
	lst.append(X)

for n in range(s):
	A = lst[0][n]
	for m in range(s):
		C = lst[2][m]
		B = lst[1][C-1]
		if A == B:
			cnt = cnt + 1
        
print(cnt) 
	

Answer 1

この問題では、（ほとんどネタバレですが）A_i = B_(C_i)なるものの個数、すなわちある i (1 <= i <= 10 ** 5)があって、i がAに何個あるか(・・・①)と、そのような（B中の） i を指し示すインデックスがCに何個あるか（・・・②）の掛け算Σ①×②で求まります。

まずこの要領の方針に気づかなければ、TLEを起こしてしまい前進は期待できません。

この方針で行っても、やり方を間違えるとTLEにはなると思いますので、別の工夫が必要です（①や②を求める.count()などの過程を工夫しない限り）。

C問題あたりだと、愚直に問題の指示通りにやってもTLEを起こしてうまくいかないので、やはりこうした問題文の適度な読み変えが必要になります。そうした新方針でやれば、仮にTLEを起こしても、「ここが計算量を重くしているな」というのが見えやすくなってくると思うので、そこまで自力で頑張られたら別途質問をされると良いと思います。（ただの愚直二重ループをはい直してくださいというと「やって欲しいことの丸投げ」と捉えられかねませんので）

Comments

[leaper]

返信遅れて申し訳ありません。詳しい解説ありがとうございます。一通り考えて質問しましたがやはり丸投げのように見える質問文だったように感じました。申し訳ありませんでした。ご教示していただいたことを今後に活かしていきたいと思います。

Answer 2

ループの中でリスト要素の要素にアクセスするのは遅いので、それぞれをAs、Bs、Csとかに入れておきましょう。

Bs[k]がAs[i]となる回数をそれぞれ数えて1回以上のものを集めて回数も覚えておく。
Cs[j]が回数１回以上のBsと一致するならその回数を足していく。

とすれば計算量は減ります。

Comments

[leaper]

返信遅れて申し訳ありません。詳しい解説ありがとうございます。今後の勉強の糧にさせていただきます。