最近のD問題は骨があって解けないことが多く、悲しみを感じる今日この頃です。
終わった後もどう解こうかもやもやしていたので、頑張って解いてみました。
今回のキーワードは、「無向グラフ」の「連結成分」。
ポイントは__「DFS」と「set」__！
問題へのリンク ABC 157 D - Friend Suggestions
コンテスト中の回答
無向グラフであることに気づいて、すぐさまDFSで解くことに。
各頂点からDFSを使って連結成分から友達でもブロックでもない、点をカウントしていきます。
2重ループであきらかに間に合わないだろうなと思っていたら、案の定TLEでした。
コンテスト中はこれでタイムアップ。
from collections import deque
N,M,K = map(int,input().split())
friend = [list(map(int,input().split())) for _ in range(M)]
block = [list(map(int,input().split())) for _ in range(K)]
g = [[False]*(N+1) for _ in range(N+1)]
for i in friend:
g[i[0]][i[1]] = True
g[i[1]][i[0]] = True
b = [[False]*(N+1) for _ in range(N+1)]
for i in block:
b[i[0]][i[1]] = True
b[i[1]][i[0]] = True
stack=deque()
ans = []
for i in range(1,N+1):
cnt = 0
visited = [0] * (N+1)
visited[i] = 1
stack.append(i)
while stack:
n = stack.pop()
for j in range(1,N+1):
if g[n][j] and visited[j]==0:
stack.append(j)
visited[j] = 1
if g[i][j] == False and b[i][j]==False:
cnt+=1
ans.append(cnt)
print(*ans)
コンテスト後の回答
inを使う場合は、ListよりSetのほうが早いと聞いて、使ってみることにしました。
いろいろ試行錯誤しながらTLEと戦った結果、以下でACをいただきました。
linkに連結成分をためていき、連結成分を求める。
連結成分の探索が終わったら、
連結成分の各点に関して引くべき友達関係とブロック関係の数を引くことで求めています。
from collections import deque
N,M,K = map(int,input().split())
friend = [list(map(int,input().split())) for _ in range(M)]
block = [list(map(int,input().split())) for _ in range(K)]
f = [set() for _ in range(N+1)]
b = [set() for _ in range(N+1)]
for i,j in friend:
f[i].add(j)
f[j].add(i)
for i,j in block:
b[i].add(j)
b[j].add(i)
stack=deque()
ans = [0]*(N+1)
visited = [0]*(N+1)
for i in range(1,N+1):
if visited[i]:
continue
link = {i}
visited[i] = 1
stack.append(i)
while stack:
n = stack.pop()
for j in f[n]:
if visited[j]==0:
stack.append(j)
visited[j] = 1
link.add(j)
for i in link:
ans[i] = len(link) - len(link & f[i]) - len(link & b[i]) - 1
print(*ans[1:])
工夫ポイント
探索済みを調べるのが明らかに無駄なので、探索済みは探索しないように追加しました。
if visited[i]:
continue
ある点に対してすべての頂点が友達関係であるかを確認していたところを、
その点に対して友達関係にあるものだけから探索するように変更しました。
for j in range(1,N+1):
# ↓↓↓
for j in f[n]:
連結成分のサイズから、友達関係またはブロック関係の数を引くのに一工夫しました。
ans[i] = len(link) - len(link & f[i]) - len(link & b[i]) - 1
つまり、どういうことかというと…
積集合を使うことで、連結成分の中から友達関係・ブロック関係の数を求めています。
f = [set for _ in range(10)]
b = [set for _ in range(10)]
# 連結成分
link = {1,2,4,9,10}
# 1と4・10が友達関係
f[1] = {4,10}
# 1は5・6・9をブロックしている
b[1] = {5,6,9}
print(link&f[1], link&b[1])
print(len(link&f[1]), len(link&b[1]))
{10, 4} {9}
2 1