[PYTHON] Atcoder Begginer Contst 065 --B: Trained?

Atcoder Begginer Contst 065 -- B: Trained? https://atcoder.jp/contests/abc065 The following two codes are all the same except for the variable called count. It was version1 => AC, version2 => TLE. There is a difference between bool and int, so I can't blame append on it, but I didn't think it would make a difference as much as TLE ...

version1.py

N = int(input())
A = []
for i in range(N):
  A.append(int(input())-1)
    
ans = -1
push = 0
count = [False]*N
for i in range(N+1):
 
    if A[push] == 1:
        ans = i+1
        break
        
    elif count[push]:
      break
      
    else:
        count[push] = True
        push = A[push] 
      
print(ans)

version2.py

N = int(input())
A = []
for i in range(N):
  A.append(int(input())-1)
    
ans = -1
push = 0
count = []
for i in range(N+1):
 
    if A[push] == 1:
        ans = i+1
        break
        
    elif A[push] in count:
      break
      
    else:
      count.append(push)
      push = A[push]
        
print(ans)

Or the 15th line of version 2

elif A[push] in count:
    break

May be slow ...? You need to be careful when handling the Ikansen list ... If anyone knows the cause, please comment.

Postscript

I changed the count variable from list to set and it no longer TLEs. After all, I think that xx in list was slow.

Conclusion: Use xx in set instead of xx in list.

version3.py

N = int(input())
A = []
for i in range(N):
  A.append(int(input())-1)
    
ans = -1
push = 0
count = set()
 
for i in range(N+1):
 
    if A[push] == 1:
        ans = i+1
        break
        
    elif A[push] in count:
      break
      
    else:
      count.add(push)
      push = A[push]
        
print(ans)