[99클럽 코테 스터디 10일차 TIL] 특정 거리의 도시 찾기 - 백준 18352

10일차 오늘도 BFS 문제다.

 

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

어떤 나라에는 1번부터 N번까지의 도시와 M개의 단방향 도로가 존재한다.
모든 도로의 거리는 1이다.
이 때 특정한 도시 X로부터 출발하여 도달할 수 있는 모든 도시 중에서, 
최단 거리가 정확히 K인 모든 도시들의 번호를 출력하는 프로그램을 작성하시오.
또한 출발 도시 X에서 출발 도시 X로 가는 최단 거리는 항상 0이라고 가정한다.

예를 들어 N=4, K=2, X=1일 때 다음과 같이 그래프가 구성되어 있다고 가정하자.
이 때 1번 도시에서 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 2인 도시는 4번 도시 뿐이다. 
2번과 3번 도시의 경우, 최단 거리가 1이기 때문에 출력하지 않는다.

[입력]
첫째 줄에 도시의 개수 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번 도시로 이동하는 단방향 도로가 존재한다는 의미다.
단, A와 B는 서로 다른 자연수이다. (1 ≤ A, B ≤ N)

출처: https://www.acmicpc.net/problem/18352

풀이 접근

그래프에 방향성이 있기 때문에 시작 노드부터 Depth대로 순회하면 된다.

BFS 돌기전 입력과 Cache(방문 여부 확인) 변수를 정의합니다.

let nmkx = readLine()!.split(separator: " ").map { Int("\($0)")! }
var graph: [Int:[Int]] = [:]
for _ in 0..<nmkx[1] {
    let road = readLine()!.split(separator: " ").map { Int("\($0)")! }
    if let _ = graph[road[0]] {
        graph[road[0]]!.append(road[1])
    } else {
        graph[road[0]] = [road[1]]
    }
}
var cache: [Bool] = Array(repeating: false, count: nmkx[0] + 1)
var q: [Int] = [nmkx[3]]
var count = 0

제출하고나서 시작 노드의 Cache 값을 저장하지 않았다는 걸 알았지만,,, 통과는 됩니다.

cache[nmkx[3]] = true

 

이후 BFS로 순회합니다.

while !q.isEmpty {
    var temp: [Int] = []
    count += 1
    if count > nmkx[2] { break }
    for num in q where !cache[num] {
        cache[num] = true
        temp.append(contentsOf: graph[num] ?? [])
    }
    if count == nmkx[2] {
        let tempSet = Array(Set(temp).filter { !cache[$0] })
        print(tempSet.count == 0 ? -1 : tempSet.sorted(by: <).map(String.init).joined(separator: "\n"))
        break
    }
    q = temp
}

루프를 통과한 후 K값이 넘어간 경우와 K값이 되기 전 더 이상 탐색할 노드가 없어 반복이 종료된 경우에 -1을 출력해줍니다.

if count != nmkx[2] {
    print(-1)
}

최종 코드 - Swift

let nmkx = readLine()!.split(separator: " ").map { Int("\($0)")! }
var graph: [Int:[Int]] = [:]
for _ in 0..<nmkx[1] {
    let road = readLine()!.split(separator: " ").map { Int("\($0)")! }
    if let _ = graph[road[0]] {
        graph[road[0]]!.append(road[1])
    } else {
        graph[road[0]] = [road[1]]
    }
}
var cache: [Bool] = Array(repeating: false, count: nmkx[0] + 1)
var q: [Int] = [nmkx[3]]
var count = 0
while !q.isEmpty {
    var temp: [Int] = []
    count += 1
    if count > nmkx[2] { break }
    for num in q where !cache[num] {
        cache[num] = true
        temp.append(contentsOf: graph[num] ?? [])
    }
    if count == nmkx[2] {
        let tempSet = Array(Set(temp).filter { !cache[$0] })
        print(tempSet.count == 0 ? -1 : tempSet.sorted(by: <).map(String.init).joined(separator: "\n"))
        break
    }
    q = temp
}
if count != nmkx[2] {
    print(-1)
}

최종 코드 - Python3

로직은 비슷하지만 시간 초과가 발생해서

좀 더 빠른 입력 함수와 deque를 사용했다.

from collections import deque
import sys
sysInput = sys.stdin.readline

n,m,k,x = map(int,sysInput().split())
graph = [[] for _ in range(n + 1)]
for _ in range(m):
  s,e = map(int,sysInput().split())
  graph[s].append(e)

cache = [-1] * (n + 1)
cache[x] = 0
q = deque([x])
result = []

while q:
  num = q.popleft()
  if cache[num] == k:
    result.append(num)
    continue
        
  for next in graph[num]:
    if cache[next] == -1:
      cache[next] = cache[num] + 1
      q.append(next)

print("\n".join(map(str, sorted(result))) if result else -1)

Conclusion

DFS보다는 BFS가 개념적으로 발견하기 더 쉬운 것 같습니다.