<백준 18352번-python> 특정 거리의 도시 찾기

풀이 시간 : 30분

시간 제한 : 2초

메모리 제한 : 256MB

출처 : 백준-18352번 https://www.acmicpc.net/problem/18352

입력

첫째 줄에 도시의 개수 N, 도로의 개수 M, 거리 정보 K, 출발 도시의 번호 X가 주어진다. (2 ≤ N ≤ 300,000, 1 ≤ M ≤ 1,000,000, 1 ≤ K ≤ 300,000, 1 ≤ X ≤ N) 둘째 줄부터 M개의 줄에 걸쳐서 두 개의 자연수 A, B가 공백을 기준으로 구분되어 주어진다. 이는 A번 도시에서 B번 도시로 이동하는 단방향 도로가 존재한다는 의미다. (1 ≤ A, B ≤ N) 단, A와 B는 서로 다른 자연수이다.

출력

X로부터 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 K인 모든 도시의 번호를 한 줄에 하나씩 오름차순으로 출력한다.

이 때 도달할 수 있는 도시 중에서, 최단 거리가 K인 도시가 하나도 존재하지 않으면 -1을 출력한다.

내가 써보는 문제 풀이

그래프가 주어지고, 최단거리를 찾는 문제이기 때문에 DFS나 BFS를 사용하면 된다. 특히 거리가 1로 고정되어 있기 때문에 BFS를 사용하는게 좋아 보인다.

그리고 이 문제는 입력량이 최대 1,000,001개로 상당히 많기 때문에 sys모듈의 stdin.readline()을 사용했다.

또한 bfs를 구현하는데 있어서 큐를 사용하기 위해 collections 모듈의 deque를 사용한다.

먼저 전체 입력을 받는 부분이다.

N, M, K, X = map(int, input().rstrip().split())
street_map = [[] for _ in range(N+1)]
for _ in range(M):
    A, B = map(int, stdin.readline().rstrip().split())
    street_map[A].append(B)

street_map은 도시와 거리 정보를 인접 리스트 형태로 구현한 2차원 리스트다. 그래프를 인접 행렬로 구현할 경우 도시 개수가 최대 300,000개 이기 때문에 그의 제곱인 90,000,000,000 = 900억 개의 원소를 표현해야 한다. 이미 메모리 초과.

found_cities = [-1] * (N + 1)
found_cities[X] = 0

found_cities는 방문한 도시를 저장하기 위한 리스트로 각 인덱스 번호가 도시 번호가 되며, 원소 값은 거리 값이 된다.

이제 마지막으로 bfs를 구현하면 된다.

queue = deque([X])
while queue:
    city = queue.popleft()
    for next_city in street_map[city]:
        if found_cities[next_city] == -1:
            found_cities[next_city] = found_cities[city] + 1
            queue.append(next_city)

핵심은 모든 거리 길이가 1이기 때문에 다음 방문할 도시와의 거리가 현재 도시의 +1인 점을 이용하는 것이다.

처음에 이런 정보를 활용하지 못해 따로 거리들을 기록하면서 딕셔너리를 만들어 줬는데, 이러니까 메모리 초과가 많이 났다.

++ 추천하는 추가 케이스

# 5가 나오는 지 확인할 수 있음
5 5 2 1
1 2
1 4
2 3
3 5
4 5

# 1이 나오는 지 확인할 수 있음
4 4 3 1
1 2
1 3
2 3
2 4
4 1

# 무한 루프 확인
3 3 4 1
1 2
2 3
3 1

전체 소스 코드 → https://github.com/on1ystar/algorithm-problem-solving-codes/blob/master/DFS%26BFS/특정 거리의 도시 찾기.py