14621번: 나만 안되는 연애

https://www.acmicpc.net/problem/14621

14621번: 나만 안되는 연애

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.Comparator;
import java.util.StringTokenizer;
import java.util.Vector;

public class java_14621 {
    final static boolean MAN = true;
    final static boolean WOMAN = false;
    final static int INF = 987654321;

    static Vector<Pair<Integer, Pair<Integer, Integer>>> edges;

    static int N, M;

    static class DisjointSet {
        Vector<Integer> parents;
        Vector<Integer> rank;

        public DisjointSet(int n) {
            parents = new Vector<>();
            rank = new Vector<>();

            for (int i = 0; i < n; i++) {
                parents.add(i);
                rank.add(i);
            }
        }

        public int find(int u) {
            if (u == parents.get(u)) {
                return u;
            }

            int ret = find(parents.get(u));
            parents.set(u, ret);
            return ret;
        }

        public void merge(int u, int v) {
            u = find(u);
            v = find(v);

            // 같은 집합에 이미 속해있음
            if (u == v) {
                return;
            }

            if (u > v) {
                int tmp = v;
                v = u;
                u = tmp;
            }

            // u의 부모를 v로 설정
            parents.set(u, v);

            // 랭크 설정
            if (rank.get(u) == rank.get(v)) {
                rank.set(v, rank.get(v) + 1);
            }
        }
    }

    static class Pair<A, B> {
        A first;
        B second;

        public Pair(A first, B second) {
            this.first = first;
            this.second = second;
        }

    }

    public static long kruskal() {
        long ret = 0;
        int count = 0;

        // cost 기반 정렬
        edges.sort(new Comparator<Pair<Integer, Pair<Integer, Integer>>>() {

            @Override
            public int compare(Pair<Integer, Pair<Integer, Integer>> o1,
                    Pair<Integer, Pair<Integer, Integer>> o2) {
                // TODO Auto-generated method stub
                return o1.first - o2.first;
            }
        });

        DisjointSet sets = new DisjointSet(N);

        for (int i = 0; i < edges.size(); i++) {
            int cost = edges.get(i).first;
            int u = edges.get(i).second.first;
            int v = edges.get(i).second.second;

            if (sets.find(u) == sets.find(v)) {
                continue;
            }

            sets.merge(u, v);

            ret += cost;

            count++;
        }

        if (count < N - 1) {
            return -1;
        }

        return ret;

    }

    public static void main(String[] args) throws IOException {
        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());

        edges = new Vector<>();

        char[] schoolType = br.readLine().replace(" ", "").toCharArray();

        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());

            int from = Integer.parseInt(st.nextToken()) - 1;
            boolean fromType = true;
            if (schoolType[from] == 'M') {
                fromType = true;
            } else {
                fromType = false;
            }

            int to = Integer.parseInt(st.nextToken()) - 1;
            boolean toType = true;
            if (schoolType[to] == 'M') {
                toType = true;
            } else {
                toType = false;
            }

            int cost = Integer.parseInt(st.nextToken());

            // 남초-남초 or 여초-여초 : 스킵
            if (fromType == toType) {
                continue;
            }

            edges.add(new Pair<Integer, Pair<Integer, Integer>>(cost, new Pair<Integer, Integer>(from, to)));
            edges.add(new Pair<Integer, Pair<Integer, Integer>>(cost, new Pair<Integer, Integer>(to, from)));

        }
        long ret = kruskal();

        bw.write(ret + "\n");

        br.close();
        bw.close();
    }
}