아이디어

• 문제의 2번, 3번 특징에서 MST 문제임을 알 수 있다.
• 문제의 1번에서, 남초-남초 간선과 여초-여초 간선은 제외해야함을 알 수 있다.
• MST를 구하기 위해 Kruskal 알고리즘을 사용했다.
  • $N \le 1\ 000$이므로 Prim을 써도 상관 없을 것이다.
• 간선 개수가 $N-1$개 미만이거나 입력이 연결그래프가 아니라면 MST 연결이 불가능하므로 -1을 출력해야 한다.

코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.StringTokenizer;

public class Main {
	static BufferedReader rd = new BufferedReader(new InputStreamReader(System.in));
	static StringTokenizer tk = null;

	static int N, M, u, v, d, ans;
	static boolean[] namcho;
	static List<Edge> edgeList = new ArrayList<>();
	static int[] parent;
	
	public static void main(String[] args) throws Exception {
		tk = new StringTokenizer(rd.readLine());
		N = Integer.parseInt(tk.nextToken());
		M = Integer.parseInt(tk.nextToken());

		namcho = new boolean[N+1];
		tk = new StringTokenizer(rd.readLine());
		for (int i=1; i<=N; i++) {
			namcho[i] = tk.nextToken().charAt(0) == 'M';
		}
		
		for (int i=0; i<M; i++) {
			tk = new StringTokenizer(rd.readLine());
			u = Integer.parseInt(tk.nextToken());
			v = Integer.parseInt(tk.nextToken());
			d = Integer.parseInt(tk.nextToken());
			
			if (namcho[u] == namcho[v]) continue;
			edgeList.add(new Edge(u, v, d));
		}
		
		// make-set
		parent = new int[N+1];
		for (int i=1; i<=N; i++) parent[i] = i;
		
		// kruskal
		edgeList.sort((e1, e2) -> Integer.compare(e1.w, e2.w));
		
		int cnt = 0;
		int idx = 0;
		int size = edgeList.size();
		while (cnt < N - 1) {
			if (idx >= size) {	// N-1개의 간선 연결 불가능
				System.out.println(-1);
				return;
			}
			Edge e = edgeList.get(idx++);
			if (union(e.v1, e.v2)) {
				cnt++;
				ans += e.w;
			}
		}
		
		// 연결 그래프인지 확인
		int p1 = findSet(1);
		for (int i=2; i<=N; i++) {
			// 같은 집합이 아니라면 분리그래프임
			if (findSet(i) != p1) {
				System.out.println(-1);
				return;
			}
		}
		
		System.out.println(ans);
	}
	
	static class Edge {
		int v1, v2, w;
		Edge(int v1, int v2, int w) {
			this.v1 = v1;
			this.v2 = v2;
			this.w = w;
		}
	}
	
	static int findSet(int x) {
		if (parent[x] == x) return x;
		return parent[x] = findSet(parent[x]);
	}
	
	static boolean union(int x, int y) {
		int px = findSet(x);
		int py = findSet(y);
		if (px == py) return false;
		parent[px] = py;
		return true;
	}
}

메모리 및 시간

• 메모리: 15960 KB
• 시간: 152 ms

리뷰

• (23.09.27 기준) 입력이 연결그래프로 주어짐이 보장되지 않는데, 주어진 간선 중 남초-여초 간선이 $N-1$개 이상이나 연결 그래프가 아닐 때를 테스트케이스에서 잡지 못하는 것 같다.
  • 해당 처리 코드를 빼도 통과가 된다.