99클럽 코테 스터디 10일차 TIL

백준 - 18352번

문제 키워드

• A번 도시에서 B번 도시로 이동하는 단방향 도로가 존재
• 도시의 최단 거리를 계산하여 k 거리만큼 떨어진 도시들의 번호 출력

최단 거리 값을 구하는 게 아닌, 최단 거리에 있는 도시들의 번호를 출력하는 것이다.

그래프 탐색인 것은 알겠는데, 이미 k라는 최단 거리 값이 있어 어떤 탐색을 써야 하나 혼동이 왔다.

 

지금까지 풀었던 BFS 문제들을 보면, 모든 노드를 탐색하고 해당 노드에 시작 노드와의 거리 값을 입력했다. 그리고 최단 거리 값을 출력했다.

풀이

위와 같은 방법으로 접근하고, 주어진 최단 거리 값이 같은 노드 즉, 도시 번호를 출력하면 된다.

• BFS를 사용하여 출발 도시에서부터 각 도시의 거리를 계산한다.
• 주어진 k 와 같다면 해당 도시 번호를 리스트에 저장한다.
• 빈 리스트가 반환되면 -1을 출력한다.

import java.util.*;

public class Main {

    static ArrayList<Integer>[] graph;
    static int[] distance;
    static boolean[] visited;
    static int k;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int n = sc.nextInt();
        int m = sc.nextInt();
        k = sc.nextInt();
        int x = sc.nextInt();
        
        // 그래프 초기화
        graph = new ArrayList[n + 1];
        distance = new int[n + 1];
        visited = new boolean[n + 1];
        for (int i = 1; i <= n; i++) {
            graph[i] = new ArrayList<>();
        }
        
        // 단방향 그래프
        for (int i = 0; i < m; i++) {
            int a = sc.nextInt();
            int b = sc.nextInt();
            graph[a].add(b);
        }

        ArrayList<Integer> results = bfs(x);

        if (!results.isEmpty()) {
            // 최단 거리에 있는 도시 오름차순 정렬
            Collections.sort(results);
            for (int result : results) {
                System.out.println(result);
            }
        } else {
            System.out.println(-1);
        }
        sc.close();
    }

    static ArrayList<Integer> bfs(int start) {
        ArrayList<Integer> results = new ArrayList<>();
        Queue<Integer> queue = new LinkedList<>();
        queue.add(start);
        visited[start] = true;

        while (!queue.isEmpty()) {
            int tmp = queue.remove();

            for (int node : graph[tmp]) {
                if (!visited[node]) {
                    visited[node] = true;
                    distance[node] = distance[tmp] + 1;

                    queue.add(node);
                    if (distance[node] == k) {
                        results.add(node);
                    }
                }
            }
        }
        return results;
    }
}

문제는 통과되었지만 개선할 부분들이 보인다.

 

1. 불필요한 노드 방문

최단 거리가 k와 같은 도시들만 리스트에 저장한다. k 이상의 거리에 있는 도시들을 탐색할 필요는 없을 것 같다.

bfs 함수 while문의 queue.add() 부분을 보면 조건이 없어 k 보다 긴 거리값까지 탐색하게 된다.

k 거리보다 짧은 거리의 도시들만 탐색할 수 있도록 조건을 추가하자.

 

2. visited 체크와 distance 계산 간소화

visited []는 노드 방문 여부를 확인하는 객체이고, distance []는 각 도시가 출발 도시에서부터 떨어진 거리값을 넣는 객체이다. 노드를 방문했다면 distance의 값이 주어질 것이고, 방문하지 않았다면 초기값 그대로 있을 것이다.

두 객체의 역할이 비슷해서 간소화해볼 수 있지 않을까 생각된다.

 

3. Stream을 사용하고 싶다.

main 함수의 마지막 부분인 sort 와 결괏값 출력을 stream으로 표현해보고자 한다.

개선 코드

import java.util.*;

public class Main {

    static ArrayList<Integer>[] graph;
    static int[] distance;
    static int k;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int n = sc.nextInt();
        int m = sc.nextInt();
        k = sc.nextInt();
        int x = sc.nextInt();

        graph = new ArrayList[n + 1];
        distance = new int[n + 1];
        Arrays.fill(distance, -1); // 거리 배열 초기값을 -1로 설정
        for (int i = 1; i <= n; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int i = 0; i < m; i++) {
            int a = sc.nextInt();
            int b = sc.nextInt();
            graph[a].add(b);
        }

        ArrayList<Integer> results = bfs(x);

        if (!results.isEmpty()) {
            results.stream()
                    .sorted()
                    .forEach(System.out::println);
        } else {
            System.out.println(-1);
        }
        sc.close();
    }

    static ArrayList<Integer> bfs(int start) {
        ArrayList<Integer> results = new ArrayList<>();
        Queue<Integer> queue = new LinkedList<>();
        queue.add(start);
        distance[start] = 0;

        while (!queue.isEmpty()) {
            int tmp = queue.remove();

            for (int node : graph[tmp]) {
                if (distance[node] == -1) { // 방문하지 않은 노드만 추가
                    distance[node] = distance[tmp] + 1;

                    if (distance[node] == k) {
                        results.add(node);
                    }

                    // 탐색 범위 축소 -> 불필요한 노드 방문 줄임
                    if (distance[node] < k) {
                        queue.add(node);
                    }
                }
            }
        }

        return results;
    }
}