본문 바로가기

Algorithm/SWEA

[SWEA-1251][MST/Prim] 하나로 - Java

문제를 딱 보면

아~~ 얘는 MST를 만들어야 하는구나~~

하는게 딱 느껴집니다.

 

MST를 만드는 알고리즘 중에서 Prim's 알고리즘을 사용해봤습니다.

import java.io.*;
import java.util.*;

public class Solution {
	static class Node{
		int to;
		long weight;
		Node(int to, long weight){
			this.to = to; this.weight = weight;
		}
	}
	
	public long Prim(List<List<Node>> graph) {
		boolean[] v = new boolean[graph.size()];
		long[] minEdge = new long[graph.size()];
		Arrays.fill(minEdge, Long.MAX_VALUE);
		
		int minVertex;
		long min, result = 0;
		minEdge[0] = 0;
		
		for(int c = 0; c < graph.size(); c++) {
			min = Long.MAX_VALUE;
			minVertex = 0;
			
			for(int i = 0; i < graph.size(); i++) {
				if(!v[i] && minEdge[i] < min) {
					min = minEdge[i];
					minVertex = i;
				}
			}
			result += min;
			v[minVertex] = true;
			
			for(int i = 0; i < graph.get(minVertex).size(); i++) {
				Node next = graph.get(minVertex).get(i);
				if(!v[next.to] && next.weight < minEdge[next.to]) {
					minEdge[next.to] = next.weight;
				}
			}
		}
		return result;
	}
	
	public static void main(String[] args) throws IOException {
		Solution s = new Solution();
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		StringTokenizer st;
		
		int T = Integer.parseInt(br.readLine());
		
		for(int tc = 1; tc <= T; tc++ ) {
			int N = Integer.parseInt(br.readLine());
			st = new StringTokenizer(br.readLine(), " ");
			long[] x = new long[N];
			long[] y = new long[N];
			List<List<Node>> graph = new ArrayList<>();
			for(int i = 0; i < N; i++) graph.add(new ArrayList<>());
			
			for(int i = 0 ; i < N; i++) x[i] = Long.parseLong(st.nextToken());
			st = new StringTokenizer(br.readLine(), " ");
			for(int i = 0 ; i < N; i++) y[i] = Long.parseLong(st.nextToken());
			
			double E = Double.parseDouble(br.readLine());

			// 그래프 만들기
			for(int i = 0; i < N; i++) {
				for(int j = 0; j < N; j++) {
					graph.get(i).add(new Node(j, (x[i] - x[j]) * (x[i] - x[j]) + (y[i] - y[j]) * (y[i] - y[j])));
				}
			}
			bw.write("#" + tc + " " + Math.round(E * s.Prim(graph)) + "\n");
		}
		
		bw.flush();
		bw.close();
	}
}

이 문제는 들어오는 인풋으로 그래프를 만드는게 좀 귀찮습니다.

to, weight 값을 가진 노드를 만들어서 리스트를 만들었습니다. 

 

그래프만 잘 만들면 프림 알고리즘을 그대로 가져다 쓰면 됩니다.

 

프림 알고리즘은 다익스트라랑 생긴게 비슷합니다. 마찬가지로 minEdge[]의 값을 우선 순위 큐에 넣어서 더 빨리 갖고 가장 작은 값을 가져올 수도 있습니다.

MST도 조만간 정리해서 글을 올려야겠습니다...

 

나올 수 있는 최댓값을 잘 생각해보고 자료형을 씁시다!!!

 

감사합니다~!