백준 5567 결혼식

https://www.acmicpc.net/problem/5567

문제를 보고 상근이의 친구와 그 친구의 친구를 찾아야하니깐 상근이와 거리가 1,2 인 사람들을 찾아야겠다고 생각하였습니다.
상근이 학번은 1 이므로 처음 큐에 1을 넣고 시작하고, 그래프를 통해 방문 가능한 (아직 방문하지 않은)친구들과의 거리를 계산하고 큐에 새로운 친구들을 넣어줍니다. 이러한 과정을 큐에 아무런 원소도 없을때까지 반복하면 visited 배열에 상근이와 친구들의 거리를 알 수 있습니다. 저는 상근이를 방문했다는 의미로 visited[1] =1 이기 때문에 visited 배열이 2와 3 사이인 경우를 세주면 답이 나옵니다.

from collections import deque
n = int(input())
m = int(input())
visited = [ 0 for _ in range(n+1)]
graph = [[] for _ in range(n+1)]

for _ in range(m):
    a,b = map(int,input().split())
    graph[a].append(b)
    graph[b].append(a)

def bfs():
    ans =0
    q= deque()
    q.append(1)
    visited[1] =1 
    while q:
        x = q.popleft()
        for nx in graph[x]:
            if not visited[nx] :
                visited[nx] = visited[x]+1
                q.append(nx)
                if 2<=visited[nx]<=3:
                    ans+=1
    return ans            
print(bfs())​