1. 문제
2. 어떻게 풀까?
MST의 비용을 구할 수 있어야 풀 수 있는 문제로 진행되는 라운드 동안 하나씩 가중치가 낮은 간선을 제거하면서 MST의 비용을 구하는 문제이다.
잊고 있던 MST 비용 구하는 개념을 다시 복습해보자
- Edge 간선을 우선순위 큐에 모두 삽입한다
- 가중치가 낮은(또는 높은)순으로 우선순위 큐에서 빼낸다
- Edge를 이루는 두 정점을 이을 시 싸이클이 형성되는지 확인한다 (find 연산)
- 싸이클이 형성되지 않으면 해당 Edge를 MST를 이루는 간선으로 추가한다.
- 간선의 개수가 N-1개가 되면 비용을 리턴한다. N-1개가 되지 못하면 MST 그래프를 만들 수 없다
먼저 MST를 구하는 코드를 짠다.
public static int MST(int k, PriorityQueue<Edge> pq){
int count = 0;
int sum = 0;
// parent 배열 자기 자신으로 초기화
for (int i = 0; i < N+1; i++) {
parent[i] = i;
}
// k개만큼 큐에서 빼주기
for (int i = 0; i < k; i++)
pq.poll();
while(!pq.isEmpty()){
Edge cur = pq.poll();
// 같은 그룹에 속하는지 확인
int a = find(cur.a);
int b = find(cur.b);
if(a!=b){
count++;
sum += cur.w; // MST를 이루는 비용 증가
// 그룹 연결
parent[b] = a;
}
}
return count == N-1? sum : 0;
}
하지만 라운드마다 MST를 형성할 수 있는지 다시 확인해야한다. 이때 고민이 되었다. 우선순위 큐를 어떻게 재활용할 수 있을까....
그냥 간선정보를 입력 받을 때 Edge를 저장하는 List를 만들어서 저장하고 `MST()`를 호출할 때마다 새로운 Que를 만들어 List에 들어있는 간선들을 넣어줬다. 이렇게 하고 정답이 되긴했지만 맘에 안들어서 다른 사람들의 코드를 찾아봤다가 새로운 사실을 알았다
아래와 같이 `PriorityQueue`를 생성할 때 데이터를 담아둔 우선순위 큐를 넣어주면 안에 들어있는 데이터와 사이즈를 그대로 복사하여 새로운 큐를 만들어준다! 굳이 새로운 큐를 만들고, for문 통해 데이터를 하나하나 넣어주지 않아도 된다😆
MST(k, new PriorityQueue<>(que));public static int MST(int k, PriorityQueue<Edge> pq){
int count = 0;
int sum = 0;
// parent 배열 자기 자신으로 초기화
for (int i = 0; i < N+1; i++) {
parent[i] = i;
}
... 생략
}
위의 방식을 사용하여 완성된 코드를 아래와 같다.
전체 코드
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.List;
import java.util.PriorityQueue;
import java.util.StringTokenizer;
public class Main {
public static int N, M, K;
public static PriorityQueue<Edge> que;
public static int [] parent;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
StringTokenizer st = new StringTokenizer(br.readLine());
N = Integer.parseInt(st.nextToken()); // 그래프 정점의 개수
M = Integer.parseInt(st.nextToken()); // 그래프 간선의 개수
K = Integer.parseInt(st.nextToken()); // 턴의 수
parent = new int[1001];
// 간선 정보 입력
que = new PriorityQueue<>();
for (int i = 1; i <= M; i++) {
st = new StringTokenizer(br.readLine());
int a = Integer.parseInt(st.nextToken());
int b = Integer.parseInt(st.nextToken());
que.add(new Edge(a, b, i));
}
boolean isEnd = false;
for (int k = 0; k < K; k++) {
if (isEnd) { // 앞 라운드에서 0이 나왔으면
bw.write("0 ");
continue;
}
int sum = MST(k, new PriorityQueue<>(que));
if (sum == 0)
isEnd = true;
bw.write(sum + " ");
}
bw.flush();
}
public static int MST(int k, PriorityQueue<Edge> pq){
int count = 0;
int sum = 0;
// parent 배열 자기 자신으로 초기화
for (int i = 0; i < N+1; i++) {
parent[i] = i;
}
// k개만큼 큐에서 빼주기
for (int i = 0; i < k; i++)
pq.poll();
while(!pq.isEmpty()){
Edge cur = pq.poll();
// 같은 그룹에 속하는지 확인
int a = find(cur.a);
int b = find(cur.b);
if(a!=b){
count++;
sum += cur.w; // MST를 이루는 비용 증가
// 그룹 연결
parent[b] = a;
}
}
return count == N-1? sum : 0;
}
public static int find(int a){
if(parent[a] == a)
return a;
return find(parent[a]);
}
}
class Edge implements Comparable<Edge>{
int a;
int b;
int w;
public Edge(int a, int b, int w) {
this.a = a;
this.b = b;
this.w = w;
}
@Override
public int compareTo(Edge o) {
// 가중치가 낮은 것을 우선순위로
return this.w - o.w;
}
}
