Atcoder ABC157 D問題 高速化の方法とWA回避

Question

発生している問題・実現したいこと

Atcoder157 ABC157 D問題にコードを提出したところ、28件中TLE11件、WA1件となってしまいました。
特にlargeとあるファイルは全てTLEでした。
https://atcoder.jp/contests/abc157/tasks/abc157_d

どこが原因でWAとなるのか、またどこを修正すればTLEとならないのかを教えていただきたいです。
よろしくお願い致します。

該当のソースコード

n,m,k = map(int,input().split())
ab = [[] for i in range(n)]
for i in range(m):
    a,b = map(int,input().split())
    ab[a-1].append(b-1)
    ab[b-1].append(a-1)
cd = [[] for i in range(n)]
for i in range(k):
    c,d = map(int,input().split())
    cd[c-1].append(d-1)
    cd[d-1].append(c-1)
    
ans = [0 for i in range(n)]
data = [0 for i in range(n)] #見たら1
for i in range(n): #自分がi
    if(data[i] != 0): #もう見てたら必要なし
        break
    data[i] = 1 #iの人を見る
    data2 = ab[i] + [i] #iが入るグラフを作る
    for j in ab[i]:
        data[j] = 1
    new_friends = [ab[i]]
    for fris in new_friends:
        new = [] #次の深さのリスト
        for fri in fris: #この層の友人を１人選ぶ
            for frfr in ab[fri]: #その友人の友人
                if(data[frfr] == 0): #見てなかったら加える
                    data[frfr] = 1 
                    new.append(frfr) #次の深さのリストに足す
                    data2.append(frfr) #グラフに足す
        if(len(new)==0):
            break
        else:
            new_friends.append(new)

    sam = len(data2)
    for j in data2:
        count = sam
        for ng in cd[j]:
            if(ng in data2):
                count -= 1 #ngがいたら除く
        ans[j] = count - len(ab[j]) -1 #自分と友達を除く


print(*ans)

参考にしたファイル

https://atcoder.jp/contests/abc157/submissions/14492334
こちらの方のコードでは全てACでした。

試したこと

PyPyで提出しても結果は同じでした。
問題中のファイルは全てうまく行きました。

他に試した入力は以下などです。
####試行
入力
4 3 1
1 2
2 3
3 4
1 4
出力
1 1 1 1

入力
2 1 0
1 2
出力
0 0

入力
3 3 0
1 2
2 3
3 1
出力
0 0 0

Answer 1

    if(data[i] != 0): #もう見てたら必要なし
        break

breakではなくcontinueでしょう。

            if(ng in data2):
                count -= 1 #ngがいたら除く

リストからデータを探すinは遅いです。それが2*k回行われるのでここだけで相当な時間です。

Comments

[shake9]

ありがとうございます！
continueの部分はうっかりしていました。
inを用いる代わりに、setの積集合を&でとって解決できました。