Today I Learned

오늘은 다익스트라 알고리즘을 활용한 코테 문제를 하나 풀었다.

영향권

문제

1~n까지의 도시가 그래프로 연결되어 있다. 이차원 배열로 {{1,6,10}, {2,3,8}...}처럼 도시번호1, 도시번호2, 거리의 꼴로 각 도시의 연결 정보가 주어진다. 도시의 수 n, 문제가 발생한 도시 배열이 주어진다. 그 도시들로 부터 거리가 K이내면 문제에 휘말린다.
문제에 휘말린 도시만 오름차순으로 배열에 담아 return해라.

풀이과정

• 다익스트라 알고리즘을 사용한다.

1. 우선 노드 클래스를 Comparable을 구현해 하나 생성하고 필드에 가중치와 정점을 int로 넣는다.

2. edge 배열을 이용해 이차원 리스트로 그래프를 양방향으로 구한다.

3. answer는 트리셋으로 만든다. 자동으로 오름차순 정렬도 되고, 중복도 걸러주니 좋다.

4. 다익스트라는 우선순위 큐를 활용해 구현한다. 한바퀴 돌면 다음 distance 배열에 k 이하 값들만 answer에 넣어준다.

5. 다익스트라를 문제가 발생한 도시마다 돌려서 answer에 휘말린 도시들을 넣는다.

6. 문제가 발생한 도시가 답에 포함될 수 있으므로 마지막으로 제거해주고 set을 배열로 바꿔서 반환한다.

Java 코드

import java.util.*;
public class Solution {
  List<List<Node>> graph;
  int[] distance;
  int K, N;
  Set<Integer> answer;

  public void dijkstra(int index){
    Queue<Node> pq = new PriorityQueue<>();
    distance = new int[N+1];
    Arrays.fill(distance,Integer.MAX_VALUE);
    distance[index] = 0;
    pq.offer(new Node(index,0));

    while (!pq.isEmpty()){
      Node node = pq.poll();
      int nv = node.v;
      int nw = node.w;

      if (nw > distance[nv]) continue;

      for (Node linked : graph.get(nv)){
        int lv = linked.v;
        int lw = linked.w;
        if (nw+lw < distance[lv]){
          distance[lv] = nw+lw;
          pq.offer(new Node(lv, distance[lv]));
        }
      }
    }
    for (int i=1; i<=N; i++){
      int ans = distance[i];
      if (ans!=index && ans<=K){
        answer.add(i);
      }
    }

  }

  public int[] solution(int n, int k, int[] city, int[][] edges){
    answer = new TreeSet<>();
    K = k;
    N = n;

    // 그래프 이차원 리스트로 생성
    graph = new ArrayList<>();
    for (int i=0; i<=n; i++){
      graph.add(new ArrayList<>());
    }

    // 양방향으로 그래프 생성
    for (int[] e : edges) {
      graph.get(e[0]).add(new Node(e[1],e[2]));
      graph.get(e[1]).add(new Node(e[0],e[2]));
    }

    for (int i = 0; i < p.length; i++) {
      int c = city[i];
      dijkstra(c);
    }

    for(int c : city){
      answer.remove(c);
    }

    //정답 반환
    int[] result = new int[answer.size()];

    int j=0;
    for (int num : answer){
      result[j++] = num;
    }
    return result;
  }

}



class Node implements Comparable<Node> {
  int v;
  int w;

  Node(int v, int w){
    this.v = v;
    this.w = w;
  }

  @Override
  public int compareTo(Node other) {
    return Integer.compare(this.w, other.w);
  }

}