백준 18352 특정 거리의 도시 찾기

문제

어떤 나라에는 1번부터 N번까지의 도시와 M개의 단방향 도로가 존재한다. 모든 도로의 거리는 1이다.

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

예를 들어 N=4, K=2, X=1일 때 다음과 같이 그래프가 구성되어 있다고 가정하자.

이 때 1번 도시에서 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 2인 도시는 4번 도시 뿐이다. 2번과 3번 도시의 경우, 최단 거리가 1이기 때문에 출력하지 않는다.

---

코드

import java.io.*;
import java.util.*;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        //N, M, K, X 입력 받기
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());
        int x = Integer.parseInt(st.nextToken()) - 1;

        //총 노드의 갯수가 3만이므로 2차원 배열이 아닌 리스트로 도로 정보 저장
        List<Integer>[] graph = new List[n];
        for (int i = 0; i < n; i++)
            graph[i] = new ArrayList<>();
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            graph[Integer.parseInt(st.nextToken()) - 1].add(Integer.parseInt(st.nextToken()) - 1);
        }

        //거리 정보 초기화(시작점에서 갈 수 있는 거리는 1, 나머지는 무한)
        int[] distance = new int[n];
        Arrays.fill(distance, Integer.MAX_VALUE);
        for (int i : graph[x])
            distance[i] = 1;

        //정답 리스트, 방문 리스트 초기화
        List<Integer> answer = new ArrayList<>();
        boolean[] visited = new boolean[n];

        //다음 방문할 노드를 저장할 우선순위큐, 거리 정보를 기준으로 가까운 노드부터 방문
        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(o -> o[1]));

        //시작점을 거리 정보 0으로 우선순위큐에 넣고, 시작점에서 갈 수 있는 노드들을 1로 넣는다
        pq.add(new int[]{x, 0});
        graph[x].forEach(o->pq.add(new int[]{o, 1}));

        while (!pq.isEmpty()) {
            int[] ptr = pq.poll();

            //방문했던 노드라면 패스
            if (visited[ptr[0]])
                continue;
            visited[ptr[0]] = true;

            //현재 거리가 k보다 크다면 나머지 노드들은 방문할 필요가 없다
            if (ptr[1] > k)
                break;

            //현재 거리가 k라면 정답 리스트에 추가
            if(ptr[1]==k)
                answer.add(ptr[0]+1);

            //현재 노드와 연결된 노드들에 대해서
            for (int i : graph[ptr[0]]) {
                if (visited[i])
                    continue;

                //기존의 거리 정보를 업데이트할 수 있으면 우선순위큐에 삽입
                if (distance[i] > ptr[1] + 1) {
                    pq.add(new int[]{i, ptr[1] + 1});
                    distance[i] = ptr[1] + 1;
                }
            }
        }

        if(answer.size()==0){
            bw.write("-1\n");
        }
        else{
            answer.sort(Comparator.naturalOrder());
            for(int i:answer)
                bw.write(i+"\n");
        }

        bw.flush();
    }
}

---

해결 과정

1. 노드가 최대 3만 개이므로 2차원 배열로 하면 사이즈가 9억이다. 따라서 리스트 형식 그래프로 저장한 후 다익스트라를 통해서 K거리인 도시들을 찾는다. 가까운 도시부터 방문하다가 K가 넘어가면 탐색할 필요가 없기 때문에 정답을 출력한다.

2. 😁