Atcoder 228 B問題について。実行時間エラーが出ます。

Question

前提・実現したいこと

Atcoder ABC228のB問題でpythonでコードを書いたのですが、TLE(実行時間制限超過エラー)が出てしまいます。

発生している問題・エラーメッセージ

TLE(実行時間制限超過エラー)

該当のソースコード

N, X = map(int, input().split())
f = input().split()

count = 0
friends = list(f)
know = [X]
 
while 1 == 1:
    X = int(friends[X-1])
    if X not in know:
        know.append(X)
    else:
        break
print(len(know))

補足情報（FW/ツールのバージョンなど）

実行時間的にどこかで無限ループに入ってしまっていると感じたのですが、while分のところをfor文に変えて試したりもしたのですが、結果が変わりませんでした。

読みにくいコードかもしれませんがよろしくお願いいたします。

Answer 1

list の in 演算子は O(n) で遅いので、set を使いましょう。

参考 : pythonの高速化のために、用途別にcollection型(list/tuple/dictionary/set)の計算量をまとめる

N, X = map(int, input().split())
f = input().split()

count = 0
friends = list(f)
know = {X}

while 1 == 1:
    X = int(friends[X-1])
    if X not in know:
        know.add(X)
    else:
        break
print(len(know))

Comments

[Dango_desu]

listとsetで実行時間に違いがあるんですね。
ありがとうございました。