[BOJ] 나만 안되는 연애(JAVA)

[14621] 나만 안되는 연애

www.acmicpc.net/problem/14621

 

Solution

• info 배열에 각 노드(대학)의 성별 정보를 저장한 다음 PQ<Edge>에 성별이 서로 다른 대학을 잇는 간선만 추가해서 최소 신장 트리를 만든다.
• answer = 최소 신장 트리를 이루는 모든 간선 비용의 합

 

 

소스코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
 
public class Main {
    static int n, m, answer;
    static int[] info, parents;
    static List<Edge> graph;
    static PriorityQueue<Edge> pq;
 
    static class Edge implements Comparable<Edge> {
        private int from;
        private int to;
        private int d;
 
        public Edge(int from, int to, int d) {
            this.from = from;
            this.to = to;
            this.d = d;
        }
 
        @Override
        public int compareTo(Edge o) {
            return this.d - o.d;
        }
    }
 
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = stoi(st.nextToken());
        m = stoi(st.nextToken());
 
        info = new int[n + 1];
        parents = new int[n + 1];
        graph = new ArrayList<>();
 
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            if (st.nextToken().charAt(0) == 'M') {
                info[i] = 0;
            } else {
                info[i] = 1;
            }
        }
 
 
        pq = new PriorityQueue<>();
        int from, to, d;
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            from = stoi(st.nextToken());
            to = stoi(st.nextToken());
            d = stoi(st.nextToken());
 
            if (info[from] != info[to]) {
                pq.offer(new Edge(from, to, d));
            }
        }
 
        for (int i = 0; i <= n; i++) {
            parents[i] = i;
        }
 
        int count = 0;
        while (!pq.isEmpty()) {
            if (count == n - 1) {
                break;
            }
            Edge edge = pq.poll();
            if (find(edge.from) != find(edge.to)) {
                union(edge.from, edge.to);
                graph.add(edge);
                count++;
            }
        }
 
        for(Edge edge : graph){
            answer+= edge.d;
        }
 
        System.out.println(count < n-1 ? -1 : answer);
    }
 
    private static void union(int a, int b) {
        a = find(a);
        b = find(b);
        if (a != b) {
            parents[b] = a;
        }
    }
 
    private static int find(int a) {
        if (parents[a] == a) {
            return a;
        }
        return parents[a] = find(parents[a]);
    }
 
    private static int stoi(String s) {
        return Integer.parseInt(s);
    }
}