백준 9344 - 도로

문제

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

리뷰

크루스칼을 이용하여 풀이할 수 있는 간단한 문제였다.

문제에 조건에 따르면 MST를 형성하며 입력에 주어진 $p-q$ 도로가 MST에
포함되는지를 도출하는 것이 관건이다. 이에 따라 크루스칼 로직을 실행하며
간선이 채택될 경우 그것이 $p-q$ 도로인지 검사하고 그 여부를 flag 변수에
저장한 후 실행 결과로 그 값을 반환하여 답을 구하였다.

로직의 시간복잡도는 크루스칼의 $O(NlogM)$으로 수렴하고 이는 $N=10,000$,
$M=20,000$인 최악의 경우에도 무난히 제한 조건 2초를 통과한다.

코드

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

import static java.lang.Integer.parseInt;

public class Main {
    static int N;
    static int[] parent;
    static PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingInt(e -> e.w));

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = parseInt(br.readLine());
        int M, p, q, u, v, w;
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();
        while (T-- > 0) {
            st = new StringTokenizer(br.readLine());
            N = parseInt(st.nextToken());
            parent = new int[N + 1];

            M = parseInt(st.nextToken());
            p = parseInt(st.nextToken());
            q = parseInt(st.nextToken());

            while (M-- > 0) {
                st = new StringTokenizer(br.readLine());
                u = parseInt(st.nextToken());
                v = parseInt(st.nextToken());
                w = parseInt(st.nextToken());

                pq.offer(new Edge(u, v, w));
            }

            sb.append(kruskal(p, q) ? "YES" : "NO").append("\n");
            pq.clear();
        }
        System.out.print(sb);
        br.close();
    }

    static boolean kruskal(int p, int q) {
        Arrays.fill(parent, -1);
        int selected = 0;

        int r1, r2;
        boolean flag = false;

        while (!pq.isEmpty() && selected < N - 1) {
            Edge e = pq.poll();

            r1 = find(e.u);
            r2 = find(e.v);

            if (r1 == r2) continue;

            if (parent[r1] < parent[r2]) {
                parent[r1] += parent[r2];
                parent[r2] = r1;
            } else {
                parent[r2] += parent[r1];
                parent[r1] = r2;
            }

            selected++;
            if (e.u == p && e.v == q || e.u == q && e.v == p)
                flag = true;
        }

        return flag;
    }

    static int find(int u) {
        if (parent[u] < 0) return u;

        return parent[u] = find(parent[u]);
    }

    static class Edge {
        int u, v, w;

        public Edge(int u, int v, int w) {
            this.u = u;
            this.v = v;
            this.w = w;
        }
    }
}

결과