문제
N 개의 마을이 있고 각 마을 마다 한 명씩 살고 있다.
X번 마을에서 파티를 하려고 하는데, 각 마을에서 X번 마을로 왕복하는 거리가 가장 큰 마을의 왕복 거리를 구하라
입력으로는 각 마을을 잇는 단방향 도로의 시작점, 끝점, 거리가 주어진다
풀이
1~N 마을에서 X번 마을로의 각각의 최적경로를 찾는다.
단방향 그래프이고 왕복을 해야하니 반대로 X번 마을에서 1~N 번 마을까지로의 최적경로도 찾는다.
각각 왕복거리를 더하고 최대 값을 출력한다.
다익스트라를 번 실행시키면 된다.
N의 범위는 1,000 이하, M의 범위는 10,000 이하이다.
시간복잡도는 이니
1억 정도 되니까 시간제한 1초에 딱 해결될 것 같다.
예시 JAVA 코드
import java.io.*;
import java.util.*;
class Main {
static class Vertex {
int v;
int c;
public Vertex(int v, int c) {
this.v = v;
this.c = c;
}
}
static int N, M, X;
static List<List<Vertex>> list = new ArrayList<>();
static int[] cost;
static int[] roundX;
static final int INF = Integer.MAX_VALUE;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
N = Integer.parseInt(st.nextToken()); // 학생 수
M = Integer.parseInt(st.nextToken()); // 도로 수
X = Integer.parseInt(st.nextToken()); // 파티 장소
for (int n=0; n<=N; n++) {
list.add(new ArrayList<Vertex>());
}
for (int m=0; m<M; m++) {
st = new StringTokenizer(br.readLine());
int from = Integer.parseInt(st.nextToken());
int to = Integer.parseInt(st.nextToken());
int weight = Integer.parseInt(st.nextToken());
list.get(from).add(new Vertex(to, weight));
}
// (N -> X) + (X -> N)
roundX = new int[N+1];
for (int n=1; n<=N; n++) {
if (n==X) continue;
dijkstra(n, X); // n->X 최단거리 roundX[n] 에 더하기
dijkstra(X, n); // X->n 최단거리 roundX[n] 에 더하기
}
Arrays.sort(roundX);
System.out.println(roundX[N]);
}
static void dijkstra(int from, int to) {
PriorityQueue<Vertex> pq = new PriorityQueue<>((o1,o2) -> o1.c - o2.c);
cost = new int[N+1];
Arrays.fill(cost, INF);
pq.add(new Vertex(from, 0));
cost[from] = 0;
while (!pq.isEmpty()) {
Vertex pv = pq.poll();
if (pv.v == to) break;
if (pv.c > cost[pv.v]) continue;
for (Vertex nv : list.get(pv.v)) {
if (cost[nv.v] > cost[pv.v] + nv.c) {
cost[nv.v] = cost[pv.v] + nv.c;
pq.add(new Vertex(nv.v, cost[nv.v]));
}
}
}
if (from != X) roundX[from] += cost[to];
else roundX[to] += cost[to];
}
}
어제 공부한 다익스트라로 한 번에 풀어져서 기분좋아서 포스팅했는데
또 하다보니 GPT가 다익스트라를 2N 번 쓰지 않고 2번만 쓰는 개선안을 알려줬다.
생각해보니까 X에서 각자 마을로 돌아가는 거리는 X 에서 다익스트라를 한 번만 돌리면 된다
X -> n 거리를 다익스트라 한 번으로 처리한 코드
import java.io.*;
import java.util.*;
class Main {
static class Vertex {
int v;
int c;
public Vertex(int v, int c) {
this.v = v;
this.c = c;
}
}
static int N, M, X;
static List<List<Vertex>> list = new ArrayList<>();
static int[] cost;
static int[] roundX;
static final int INF = Integer.MAX_VALUE;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
N = Integer.parseInt(st.nextToken()); // 학생 수
M = Integer.parseInt(st.nextToken()); // 도로 수
X = Integer.parseInt(st.nextToken()); // 파티 장소
for (int n=0; n<=N; n++) {
list.add(new ArrayList<Vertex>());
}
for (int m=0; m<M; m++) {
st = new StringTokenizer(br.readLine());
int from = Integer.parseInt(st.nextToken());
int to = Integer.parseInt(st.nextToken());
int weight = Integer.parseInt(st.nextToken());
list.get(from).add(new Vertex(to, weight));
}
// (N -> X) + (X -> N)
roundX = new int[N+1];
for (int n=1; n<=N; n++) {
if (n==X) continue;
dijkstra(n, X); // n->X 최단거리 roundX[n] 에 더하기
}
dijkstra(X, 0); // 복귀 거리 구하기
Arrays.sort(roundX);
System.out.println(roundX[N]);
}
static void dijkstra(int from, int to) {
PriorityQueue<Vertex> pq = new PriorityQueue<>((o1,o2) -> o1.c - o2.c);
cost = new int[N+1];
Arrays.fill(cost, INF);
pq.add(new Vertex(from, 0));
cost[from] = 0;
while (!pq.isEmpty()) {
Vertex pv = pq.poll();
if (pv.v == to) break;
if (pv.c > cost[pv.v]) continue;
for (Vertex nv : list.get(pv.v)) {
if (cost[nv.v] > cost[pv.v] + nv.c) {
cost[nv.v] = cost[pv.v] + nv.c;
pq.add(new Vertex(nv.v, cost[nv.v]));
}
}
}
if (from != X) {
roundX[from] += cost[X]; // x -> X 거리
} else {
for (int n=1; n<=N; n++) {
roundX[n] += cost[n]; // X -> n 거리
}
}
}
}
위의 것이 두 번째 코드의 결과이다. 시간과 메모리가 소폭 감소하였다.
그런데 각 마을에서 X 마을로 가는 최단 거리도 다익스트라 한 번으로 구할 수 있다고 한다.
그래프 뒤집기
그래프의 방향을 반대로 뒤집으면
n -> X 로 가는 간선이 X -> n 간선으로 바뀌기 때문에
X에서 다익스트라를 돌리면 각 마을에서 X 로의 거리를 구할 수 있는 것이다.
가보자!!
역방향 그래프를 사용해 다익스트라를 총 2번만 사용한 코드
import java.io.*;
import java.util.*;
class Main {
static class Vertex {
int v;
int c;
public Vertex(int v, int c) {
this.v = v;
this.c = c;
}
}
static int N, M, X;
static List<List<Vertex>> originList = new ArrayList<>();
static List<List<Vertex>> reverseList = new ArrayList<>();
static int[] cost;
static int[] roundX;
static final int INF = Integer.MAX_VALUE;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
N = Integer.parseInt(st.nextToken()); // 학생 수
M = Integer.parseInt(st.nextToken()); // 도로 수
X = Integer.parseInt(st.nextToken()); // 파티 장소
for (int n=0; n<=N; n++) {
originList.add(new ArrayList<Vertex>());
reverseList.add(new ArrayList<Vertex>());
}
for (int m=0; m<M; m++) {
st = new StringTokenizer(br.readLine());
int from = Integer.parseInt(st.nextToken());
int to = Integer.parseInt(st.nextToken());
int weight = Integer.parseInt(st.nextToken());
originList.get(from).add(new Vertex(to, weight));
reverseList.get(to).add(new Vertex(from, weight));
}
// (N -> X) + (X -> N)
roundX = new int[N+1];
dijkstra(X, reverseList); // n -> X 거리 구하기
dijkstra(X, originList); // X -> n 거리 구하기
Arrays.sort(roundX);
System.out.println(roundX[N]);
}
static void dijkstra(int from, List<List<Vertex>> list) {
PriorityQueue<Vertex> pq = new PriorityQueue<>((o1,o2) -> o1.c - o2.c);
cost = new int[N+1];
Arrays.fill(cost, INF);
pq.add(new Vertex(from, 0));
cost[from] = 0;
while (!pq.isEmpty()) {
Vertex pv = pq.poll();
if (pv.c > cost[pv.v]) continue;
for (Vertex nv : list.get(pv.v)) {
if (cost[nv.v] > cost[pv.v] + nv.c) {
cost[nv.v] = cost[pv.v] + nv.c;
pq.add(new Vertex(nv.v, cost[nv.v]));
}
}
}
for (int n=1; n<=N; n++) {
roundX[n] += cost[n];
}
}
}
와우!!
메모리와 시간이 확연히 줄어들었다!
역시 알고리즘을 풀 때 GPT에게 검토받는 것은 좋은 습관이다.
자두과자