G14621. 나만 안되는 연애

문제

제출 코드

• 사용 알고리즘 : union-find, Kruskal, MST

간선들을 우선순위 큐에 넣어 가중치가 낮은 순으로 사이클이 되지 않는 선에서 꺼내어 트리를 만든다. n개의 정점을 연결하는 최소 간선 수(n-1) 만큼 꺼내면 끝.

꺼낸 간선을 연결하면 사이클이 생기는지 확인하는 과정에서 union-find를 사용하였는데, 간선을 연결해 부모를 통일하는 과정중에 변경할 트리의 루트노드의 부모값을 변경하지 않고 하위노드의 부모값을 변경하는 실수가 있었다. union-find 문제 풀때 자주 실수하는 부분인 듯 하다.

또 하나 유의할 점은, 노드의 부모를 찾는 findParent 구현 시 구한 부모 값을 return하기 전에 현재 노드에도 구한 노드를 저장해두는 것이 좋다. 다음에 동일한 노드의 부모를 찾아야 할 경우에, 다시 상위까지 올라가서 부모를 찾는 반복을 줄일 수 있다.

package gold;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class G14621_noMoreLove {

	static int[] parent;

	static class Edge implements Comparable<Edge>{
		int from, to, weight;
		Edge(int from, int to, int weight){
			this.from = from;
			this.to = to;
			this.weight = weight;
		}
		@Override
		public int compareTo(Edge o) {
			return this.weight - o.weight;
		}
	}

	static int findParent(int a) {
		if(parent[a]!=a) parent[a] = findParent(parent[a]);
		return parent[a];
	}

	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());

		char[] gender = new char[N+1];
		parent = new int[N+1];

		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++) {
			gender[i] = (char) st.nextToken().charAt(0);
			parent[i] = i;
		}

		PriorityQueue<Edge> queue = new PriorityQueue<Edge>();
		for(int i=0; i<M; i++) {
			st = new StringTokenizer(br.readLine());
			int from = Integer.parseInt(st.nextToken());
			int to = Integer.parseInt(st.nextToken());
			int weight = Integer.parseInt(st.nextToken());

			if(gender[from] != gender[to]) queue.add(new Edge(from, to, weight));
		}


		int cnt=0, sum=0;
		boolean[] visit = new boolean[N+1];

		while(cnt<N-1 && !queue.isEmpty()) {
			Edge now = queue.poll();
			int parent_from = findParent(now.from);
			int parent_to = findParent(now.to);
			if(parent_from == parent_to) continue;

			parent[parent_to] = parent_from;
			sum += now.weight;
			cnt++;
		}
		System.out.println((cnt==N-1) ? sum : -1);
	}

}