[백준] 1854. K번째 최단 거리

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

이번에도 논리 점프로 인해 스스로 “왜 이러는데”라고 물어보면
“그냥..;;”의 반복이라서 스스로를 납득시키기 위해 생각해봤다.

 

 

➡️ 1번 도시에서 각 도시에 대한 k번째 최단경로의 소요시간을 구해야 한다.

 

문제를 처음 읽어보면,

출발점은 1번 도시에서 고정되고, 각 도시에 도착하는 (k번째) 최단 경로의 소요시간을 구하라고 한다.
그러면 자연스럽게 떠오르는 건 모든 도시에 대한 최단 경로를 구하는 다익스트라 알고리즘이다.

다익스트라는 가장 짧은 경로(1번째 최단 경로)만 구해주는데… 여기서는 k번째 최단 경로를 구해야 한다.

 

 

🫨 생각해본다...

 

 

생각난 알고리즘을 기반으로 생각해본다.

“어떻게든 다익스트라 알고리즘을 수정하면 되지 않을까”라는 생각으로 구조를 한번 써본다.

1. 각 도시에 도착하는 최단 거리 배열을 INF로 초기화
2. 우선순위 큐를 사용해서 "가장 짧은 거리 후보"부터 꺼내 탐색
3. 더 짧은 경로를 찾으면 업데이트하고 큐에 저장

 

😠 구조를 째려본다...

 

눈에 들어오는 건…

1. 각 도시에 도착하는 최단 거리 배열을 INF로 초기화
2. 우선순위 큐를 사용해서 "가장 짧은 거리 후보"부터 꺼내 탐색
3. 더 짧은 경로를 찾으면 업데이트하고 큐에 저장

다익스트라 알고리즘에서는 더 짧은 경로를 찾으면 기존의 최단 거리 값을 버리고 새로운 값으로 업데이트하기 때문에 결국 각 도시가 딱 1개의 최단 거리를 갖게 된다.

 

하지만!!! 지금 k번째 최단 경로를 구하는 과정에서는 가장 짧은 하나만 남길 필요가 없다!! 라고 생각하니까 각 도시까지 도착하는 거리를 몽땅 저장한 다음에, 정렬해서 k번째 거를 가져오면 되지않나? 라는 생각을 했다.

 

근데 또 생각해보니까 몽땅 저장하기엔 비효율적일 것 같기도 하고, 지금 필요한 거리 개수는 최대 k개다. 그것도 짧은 순으로 k개다. 그러니까 각 도시마다 도착하는 거리 목록을 최대 k개까지만 유지하면 되는 거다. 최대라고 한 이유는 문제에서 k번째 최단 거리가 없는 도시도 있다고 언급했기 때문이다.

 

(1) 리스트를 이용해보자

다익스트라 알고리즘을 기반으로 논리를 이어보면,

새로운 경로를 발견했을 때:
1. 리스트가 k개보다 적으면 그냥 저장한다.
2. 리스트가 k개라면, 리스트에 있는 거리 중에 새로운 경로보다 긴 경로가 있으면 바꿔치기한다. ⇒ 저장된 길이 중에 가장 긴 거랑 비교한다.

 

이렇게 하면 각 도시의 거리 리스트에는 짧은 거리 상위 k개만 남게 될 거라고 생각한다….

 

근데 이렇게 되면 —기존에 저장되어 있는 최단거리 값과 새로 계산한 최단거리 값과 비교해서 짧을 때만 업데이트하고 큐에 추가하던 로직을 어떻게 바꿀 것인가— 에서 멈추게 된다. (나만 그럴 수도 있다)

리스트가 k개보다 적으면 그냥 저장한다. + 큐에 추가한다.
리스트가 k개라면, 리스트에 있는 거리 중에 새로운 경로보다 긴 경로가 있으면 바꿔치기한다. + 큐에 추가한다.

어렵게 생각하지 말고! 바꾸기 전에도 그랬듯이 새로운 값을 저장하게 되면 큐에 추가한다.

 

구현해보면:

public class Main {

  static class Edge {
    int to, dist;
    public Edge(int to, int dist) {
      this.to=to; this.dist=dist;
    }
  }
  static class Load implements Comparable<Load> {
    int to, dist;
    public Load(int to, int dist) {
      this.to=to; this.dist=dist;
    }
    @Override
    public int compareTo(Load l) {
      return this.dist - l.dist;
    }
  }
  
  public static void main(String[] args) throws IOException{
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    
    StringTokenizer st = new StringTokenizer(br.readLine());
    int n = Integer.parseInt(st.nextToken()); //도시 개수 
    int m = Integer.parseInt(st.nextToken()); //도로 개수 
    int k = Integer.parseInt(st.nextToken());
    
    //방향 가중치 그래프 
    Map<Integer, List<Edge>> map = new HashMap<>();
    for(int i=0; i<m; i++) {
      st = new StringTokenizer(br.readLine());
      int a = Integer.parseInt(st.nextToken());
      int b = Integer.parseInt(st.nextToken());
      int c = Integer.parseInt(st.nextToken());
      map.computeIfAbsent(a, key->new ArrayList<>()).add(new Edge(b,c));
    }
    
    Map<Integer, List<Integer>> shorts = dijkstra(map, n, k);
    //k번째 최단거리 출력
    for(int node=1; node<=n; node++) {
      List<Integer> ways = shorts.get(node);
      if(ways == null || ways.size() < k) System.out.println(-1);
      else {
        Collections.sort(ways);
        System.out.println(ways.get(k-1));
      }
    }
  }
  
  static Map<Integer, List<Integer>> dijkstra(Map<Integer, List<Edge>> map, int n, int k) {
    Map<Integer, List<Integer>> ways = new HashMap<>();
    
    Queue<Load> pq = new PriorityQueue<>(); //짧은 거리 우선 방문 
    pq.offer(new Load(1, 0)); //1번에서 출발 
    ways.put(1, new ArrayList<>(List.of(0)));
    
    while(!pq.isEmpty()) {
      Load curr = pq.poll();
      
      if(!map.containsKey(curr.to)) continue;
      
      for(Edge next : map.get(curr.to)) {
        int newDist = curr.dist + next.dist;
        
        //현재 저장된 거리 개수가 k보다 작으면 그냥 저장 
        if(ways.get(next.to) == null || ways.get(next.to).size() < k) {
          ways.computeIfAbsent(next.to, key -> new ArrayList<>()).add(newDist);
          pq.offer(new Load(next.to, newDist));
        }
        //현재 저장된 거리 중 가장 긴 거리가 새로 찾은 거리보다 길면 바꿔치기 
        else {
          int far = Collections.max(ways.get(next.to));
          if(far > newDist) {
            int idx = ways.get(next.to).indexOf(far);
            ways.get(next.to).set(idx, newDist);
            pq.offer(new Load(next.to, newDist));
          }
        }
      }
    }
    
    return ways;
  }
  
}

작성하면서도 (아..너무 귀찮은데)라고 생각했는데, 이렇게 돌려보면 한 19% 쯤에서 시간 초과가 발생한다.👹

 

시간 초과가 나는 부분이 어딜지 또 구조를 째려보면,

else {
  int far = Collections.max(ways.get(next.to)); // O(k)
  if(far > newDist) {
    int idx = ways.get(next.to).indexOf(far); //O(k)
    ways.get(next.to).set(idx, newDist);
    pq.offer(new Load(next.to, newDist));
  }
}

역시 작성하면서도 (음……….) 했던 이 부분일 것이다.

문제 조건에서 최악의 경우 도시 수는 1000개, 도로 수는 250,000개, k는 100개 ⇒ mk는 3,000,000개다.
그럼 이 과정이 모든 도로마다 반복될거니까 결국 시간 복잡도가 O(mk)가 되는데
O(mk) = 250,000 x 100 = 25,000,000 ⇒ 2.5 * 10^7 이 되니까 시간 초과 발생 확률이 높아진다.

 

(2) 우선순위 큐(힙)로 바꿔보자

그럼 이 부분의 효율성을 빠르게 개선할 수 있는 부분을 찾기 위해 또 생각해본다.

계속해서 <저장된 거리 중 가장 긴 거리를 찾고 + 바꾸는 부분> 이 부분을 바꿀 수 있을 거같은데
— ”반복문마다 항상 저장된 값 중 가장 큰 값”을 확인하고 그 값과 바꿔야 한다.

 

정렬 상태를 유지할 수 있는 자료 구조를 리스트 대신 사용하면 좋을 것 같다.

 

...정렬 상태를 유지할 수 있는 건.. 우선 순위 큐!! 심지어 가장 큰 값을 계속 확인하고 + 그 값만 빼서 바꿔치기도 편하다(정렬도 빠르다).

 

구현해보면:

static Map<Integer, PriorityQueue<Integer>> dijkstraWithPQ(Map<Integer, List<Edge>> map, int n, int k) {
  //큰 값이 우선 => 더 오래 걸리는 거리가 앞으로 -> head값을 꺼내면 k번째 최단 거리가 됨 
  Map<Integer, PriorityQueue<Integer>> maxHeap = new HashMap<>();
  for(int node=1; node<=n; node++) {
    maxHeap.put(node, new PriorityQueue<>((o1, o2) -> o2-o1)); 
  }
  
  Queue<Load> pq = new PriorityQueue<>(); //짧은 거리 우선 방문 
  pq.offer(new Load(1, 0)); //1번에서 출발 
  maxHeap.get(1).offer(0);
  
  while(!pq.isEmpty()) {
    Load curr = pq.poll();
    
    if(!map.containsKey(curr.to)) continue;
    
    for(Edge next : map.get(curr.to)) {
      int newDist = curr.dist + next.dist;
      
      if(maxHeap.get(next.to).size() < k) {
        maxHeap.get(next.to).offer(newDist);
        pq.offer(new Load(next.to, newDist));
      }
      else {
        if(maxHeap.get(next.to).peek() > newDist) {
          maxHeap.get(next.to).poll();
          maxHeap.get(next.to).offer(newDist);
          pq.offer(new Load(next.to, newDist));
        }
      }
    }
  }
  
  return maxHeap;
}

 

여기서 효율성을 다시 확인해보면

else {
  if(maxHeap.get(next.to).peek() > newDist) { //O(1) 
    maxHeap.get(next.to).poll();
    maxHeap.get(next.to).offer(newDist); //O(logk)
    pq.offer(new Load(next.to, newDist));
  }
}

힙에서는 top 원소 확인 O(1), 원소 제거 및 추가로 인한 정렬 O(logn)이므로 여기서 시간 복잡도는 O(mlogk)로 개선할 수 있다. 

그러면 O(m*logk) = 250,000 x log100 = 250,000 x 7 = 1,750,000 ⇒ 대충 2 * 10^6

 

출력할 때도

Map<Integer, PriorityQueue<Integer>> shorts = dijkstraWithPQ(map, n, k);

//k번째 최단거리 출력
for(int node=1; node<=n; node++) {
  System.out.println((shorts.get(node).size() < k) ? -1 : shorts.get(node).peek());
}

이렇게 크기만 확인하고 peek()으로 헤드값을 출력해주면 되니 편하다.