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

막 어려운 문제라기보다는,

생각의 전환을 잘 하고 싶어서 글을 적게 된 문제이다.

 

 

문제

https://school.programmers.co.kr/learn/courses/30/lessons/132266?language=java

프로그래머스

 

 

내가 푼 풀이

처음에는 dfs + dp를 이용해서 문제를 풀려고 했었다.

이전에 갔던 길이라면, 최단거리를 다시 구할 필요가 없기 때문이다.

 

이렇게 하니 답은 나왔지만, 시간 초과가 떠서 맞지 못했다.

아래는 그 dfs 코드이다.

이미 최단거리를 구했던 곳이라면 더이상 가지않고 최단거리를 반환하여 바로 구할수 있도록 했다.

private int dfs(int cur, int cnt, boolean[] visited, int dest) {
    if(cur == dest) {
        return cnt;
    }
    if(dp[cur] > 0) {
        return dp[cur] + cnt;
    }

    int minVal = Integer.MAX_VALUE;
    for (Integer next : roadList[cur]) {
        if(visited[next]) continue;

        visited[next] = true;
        minVal = Math.min(minVal, dfs(next, cnt+1, visited, dest));
        visited[next] = false;
    }

    if(minVal == Integer.MAX_VALUE) {
        dp[cur] = minVal;
    } else {
        dp[cur] = minVal - cnt;
    }
    return minVal;
}

 

그렇지만 이 코드는 결론적으로 시간초과가 났다.

방법들을 살펴보다가 다익스트라를 통해서 푼 것들을 발견했다.

 

 

 

다른 사람의 풀이

다익스트라는 한 정점에서 다른 모든 정점으로 가는 최단경로를 구하는 것이기 때문에,

이 문제는 모든 부대원이 강철부대 지점인 한 정점으로 간다는 점에서 다익스트라를 쓰기에 적합하다.

 

거꾸로 도착 지점인 강철부대 지점을 기준으로 하여 구하면 된다.

이렇게 반대로 생각하면 쉽게 풀 수 있는 문제인데, 이 생각을 하지 못했다.

이 생각을 하고 나면, 다익스트라를 알고 있다면 풀기에는 어렵지 않다.

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

class Main {
	List<Integer>[] roadList;
	int[] d;

	public int[] solution(int n, int[][] roads, int[] sources, int destination) {
		int[] answer = new int[sources.length];
		roadList = new List[n+1];
		d = new int[n+1];
		Queue<int[]> pq = new PriorityQueue<int[]>((o1, o2) -> Integer.compare(o1[1], o2[1]));

		for (int i = 0; i <= n; i++) {
			roadList[i] = new ArrayList<>();
			d[i] = Integer.MAX_VALUE;
		}

		for (int[] road : roads) {
			roadList[road[0]].add(road[1]);
			roadList[road[1]].add(road[0]);
		}

		d[destination] = 0;
		pq.add(new int[]{destination, 0});

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

			if(d[x[0]] < x[1]) continue;

			for (Integer next : roadList[x[0]]) {
				if(d[next] > (d[x[0]] + 1)) {
					pq.add(new int[]{next, d[x[0]] + 1});
					d[next] = d[x[0]] + 1;
				}
			}
		}

		for (int i = 0; i < sources.length; i++) {
			if(d[sources[i]] == Integer.MAX_VALUE) {
				answer[i] = -1;
				continue;
			}
			answer[i] = d[sources[i]];
		}

		return answer;
	}
}

 

우선순위 큐를 이용하여 쉽게 풀 수 있었다.

 

 

 

 

생각을 전환해서 여러가지 알고리즘을 이용하여 잘 풀 수 있도록

많이 문제를 접해보고, 풀 때 협소하게 생각하지 않도록 노력해야겠다.