[프로그래머스] 부대복귀

프로그래머스 : 부대복귀 (Java)

문제

링크 확인

풀이

최단 거리를 구하는 문제이다.
문제가 어려운 것 같지만 차근차근 읽어보면 생각보다 쉽게 문제를 풀 수 있다.
문제에서 주어진 조건을 정리하자면 다음과 같다.

• n은 총 지역의 수를 의미한다.
• roads는 roads[x][2]의 크기로 roads[x][0]과 roads[x][1]은 서로 이어진 경로가 있다는 것을 의미한다.
• road[x][0]에서 road[x][1]로 경로가 있다는 것은, road[x][1]에서 road[x][0]으로도 이동할 수 있다는 의미이다.
• 각 지역의 이동 비용은 1이다.
• sources는 1차원 배열로 부대원들이 각각 출발하는 지역을 의미한다.
• destination은 부대원들이 도착해야 할 목적지를 말한다.

여기서 중요한 점은 주어진 경로는 왕복이라는 것이다.
이 점을 이용하면, 여러 출발지에서 하나의 목적지로의 각 최단거리를 구하는 것이 아니라, 목적지 destination을 출발지로 보고, sources 1차원 배열의 값들을 목적지로 봐도 답은 똑같이 도출될 것 이다.

하나의 출발지에서 다른 노드들의 최단거리를 구하기 위해, 우리는 다익스트라 알고리즘을 이용할 수 있다.

우선순위 큐를 사용하여 다익스트라 알고리즘을 구현하면 출발 지역 destination에 대한 각 지역으로의 최단 거리를 알 수 있고, 여기서 필요한 지역들(sources)만 꺼내어 반환하면 된다.

다익스트라 알고리즘에 대한 설명은 여기에서 확인할 수 있다.

한가지 유의할 점은, roads 배열을 통해 경로를 표현하기 위해 인접 리스트 방식을 사용해야 한다는 점이다.
본인은 처음 문제를 풀 때 인접 행렬 방식으로 문제를 풀려고 시도했다가 메모리 초과가 떠서 실패했었다.
이후 문제를 다시 보니 n의 값이 최대 100,000 이었고, 인접 행렬 방식으로 문제를 풀게 되면 100,000 * 100,000 크기의 2차원 배열이 생성되기 때문에 메모리 초과가 발생했던 것이었다.

좀 더 효율적으로 메모리를 사용할 수 있는 인접 리스트 방식을 이용하여 문제를 풀도록 하자.

코드

Java로 작성한 전체 코드는 다음과 같다.

public int[] solution(int n, int[][] roads, int[] sources, int destination) {
    int INF = 9_999_999; // 무한을 의미한다.
    int[] answer = new int[sources.length];

    // 인접 리스트 생성
    List<Node>[] list = new ArrayList[n + 1];
    for (int i = 1; i <= n; i++) {
        list[i] = new ArrayList<>();
    }
    // 경로 정보(roads)를 통해 인접 리스트 초기화
    for (int i = 0; i < roads.length; i++) {
        int start = roads[i][0];
        int end = roads[i][1];
        list[start].add(new Node(end, 1));
        list[end].add(new Node(start, 1));
    }

    // 최단 거리 테이블 생성 및 초기화
    int[] shortest = new int[n + 1];
    for (int i = 0; i <= n; i++) {
        shortest[i] = INF;
    }

    PriorityQueue<Node> priorityQueue = new PriorityQueue<>((o1, o2) -> {
        if (o1.cost == o2.cost) {
            return o1.node - o2.node;
        }
        return o1.cost - o2.cost;
    });
    // 각 지역 방문 여부 체크 테이블
    boolean[] visited = new boolean[n + 1];
    priorityQueue.add(new Node(destination, 0));
    shortest[destination] = 0;

    while (!priorityQueue.isEmpty()) {
        Node poll = priorityQueue.poll();
        int now = poll.node;

        // 방문하지 않았다면 수행
        if (!visited[now]) {
            // 현재 노드 방문 처리
            visited[now] = true;

            int size = list[now].size();
            for (int i = 0; i < size; i++) {
                Node node = list[now].get(i);

                int pathNodeCost = node.cost;
                if (now != node.node && pathNodeCost != INF) {
                    int fullCost = poll.cost + pathNodeCost;
                    shortest[node.node] = Math.min(shortest[node.node], fullCost);
                    priorityQueue.add(new Node(node.node, fullCost));
                }
            }
        }
    }

    for (int i = 0; i < sources.length; i++) {
        int source = sources[i];
        // 문제 요구사항. 복귀가 불가능하다면(INF) -1로 반환
        int a = shortest[source] == INF ? -1 : shortest[source];
        answer[i] = a;
    }

    return answer;
}

private static class Node {
    public int node; // 지역 번호
    public int cost; // 이동 비용

    public Node(int node, int cost) {
        this.node = node;
        this.cost = cost;
    }
}