백준 1504 특정한 최단 경로 Java

알고리즘 분류 : 최단 거리, 다익스트라

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

요약

두 노드 사이에 가중치가 있는 그래프이므로 다익스트라를 이용할 수 있다. 이 때 두 정점을 통과해야 한다는 추가 조건이 있는데 두 정점을 A, B라고 하자. 이 문제는 1번 정점 -> A -> B -> N번 정점, 1번 정점 -> B -> A -> N번 정점 두가지 경우로 나눌 수 있다. 두가지 경우 모두 시작점이 1번 정점, A, B이므로 이 세 정점에서만 다른 모든 정점까지 최단 거리를 다익스트라로 구하면 풀 수 있다. 이 문제를 풀 때 경로가 없을 경우일 때 오버플로우를 유의해야 한다.

풀이

다익스트라는 한 정점에서 다른 모든 정점까지의 최단 거리를 구할 수 있다. 이를 유의하면 불필요한 계산을 줄일 수 있다. 두 정점을 거칠 때 가능한 경로는 두가지이다.

1. 1번 정점 -> A -> B -> N번 정점
2. 1번 정점 -> B -> A -> N번 정점

전체 경로를 세부적으로 나누면 다음과 같다.

1. 1번 정점 -> A
2. 1번 정점 -> B
3. A -> B
4. A -> N번 정점
5. B -> N번 정점
6. B -> A

(1, 2), (3, 4), (5, 6)은 시작점이 각각 동일하므로 다익스트라를 세 번만 사용하여 구할 수 있다.

이 때 경로가 존재하지 않을 시 Integer.MAX_VALUE를 반환하기로 하였다면 세부 경로들을 더하기 전에 각각 Integer.MAX_VALUE인지 확인을 해야 Integer.MAX_VALUE에서 값이 더해져서 오버플로우가 일어나는 상황을 막을 수 있다.

코드

public class Main {
  static int N;
  static int E;
  static int start;
  static int end;
  static List<List<Node>> edgeInfo = new ArrayList<>();
  public static void main(String[] args) throws Exception {
    getInput();

    int result = -1;
    int[] distanceFromOne = dijkstra(0);
    int[] distanceFromStart = dijkstra(start);
    int[] distanceFromEnd = dijkstra(end);

    //0 -> start -> end -> n-1
    int oneToStart = distanceFromOne[start];
    int startToEnd = distanceFromStart[end];
    int endToN = distanceFromEnd[N - 1];

    if (oneToStart != Integer.MAX_VALUE &&
      startToEnd != Integer.MAX_VALUE &&
      endToN != Integer.MAX_VALUE) result = oneToStart + startToEnd + endToN;

    //0 -> end -> start -> n-1
    int oneToEnd = distanceFromOne[end];
    int EndToStart = distanceFromEnd[start];
    int StartToN = distanceFromStart[N - 1];

    if (oneToEnd != Integer.MAX_VALUE &&
      EndToStart != Integer.MAX_VALUE &&
      StartToN != Integer.MAX_VALUE) {
      if (result == -1) result = oneToEnd + EndToStart + StartToN;
      else result = Math.min(result, oneToEnd + EndToStart + StartToN);
    }

    System.out.println(result);
  }

  static void getInput() throws Exception {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    int[] line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
    N = line[0];    E = line[1];

    for (int i = 0; i < N; i++) {
      edgeInfo.add(new ArrayList<>());
    }

    for (int i = 0; i < E; i++) {
      line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
      int edgeStart = line[0] - 1; //정점 번호는 1빼기
      int edgeEnd = line[1] - 1;
      int len = line[2];

      edgeInfo.get(edgeStart).add(new Node(edgeEnd, len));
      edgeInfo.get(edgeEnd).add(new Node(edgeStart, len));
    }

    line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
    start = line[0] - 1;
    end = line[1] -1;
  }

  static int[] dijkstra(int start) {
    //priority queue
    PriorityQueue<Node> pq = new PriorityQueue<>();

    //거리 테이블 초기화
    int[] distanceFromStart = new int[N];
    for (int i = 0; i < N; i++) {
      distanceFromStart[i] = Integer.MAX_VALUE;
    }


    distanceFromStart[start] = 0;
    pq.add(new Node(start, 0));

    while (!pq.isEmpty()) {
      Node current = pq.poll();
      int currentIndex = current.index;

      //업데이트 되지 않은 Node인지 확인
      if (current.distance > distanceFromStart[currentIndex]) continue;

      List<Node> currentNodeEdgeInfo = edgeInfo.get(currentIndex);
      for (Node n : currentNodeEdgeInfo) {
        int nearNodeIndex = n.index;
        if (distanceFromStart[n.index] > distanceFromStart[currentIndex] + n.distance) {
          distanceFromStart[n.index] = distanceFromStart[currentIndex] + n.distance;
          pq.add(new Node(nearNodeIndex, distanceFromStart[n.index]));
        }
      }
    }

    return distanceFromStart;
  }

  static class Node implements Comparable<Node> {
    int index;
    int distance;

    Node(int i, int d) {
      index = i;
      distance = d;
    }

    @Override
    public int compareTo(Node n) {
      return this.distance - n.distance;
    }
  }
}

Integer.MAX_VALUE를 더할 때 오버플로우의 처리를 어떻게 할 지는 여러 방법이 있을 수 있다. 나의 경우는 더하기 전에 세부 경로가 모두 Integer.MAX_VALUE가 아닌 경우에만 유효한 값으로 판단하였다. 또한 1238 파티와 다르게 양방향이므로 반대쪽으로 오는 경로도 처음 입력을 받을 때 추가시켜 주어야 한다.

시행착오

처음에 양방향임을 생각하지 못해 반대쪽에서 오는 경로를 추가시켜주지 않아 오답이 나왔었다.

      //정방향
      edgeInfo.get(edgeStart).add(new Node(edgeEnd, len));
      //역방향
      edgeInfo.get(edgeEnd).add(new Node(edgeStart, len));

그래프에서 방향의 여부는 중요한 만큼 문제에서 주어지면 잊지 말고 체크하여 추가해야 한다.