백준 1865 - 웜홀

문제

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

풀이

• 웜홀을 통해선 음의 가중치가 발생할 수 있다는 점에서 벨만-포드를 적용할 수 있음
• 풀이에 있어 중점이 되는 부분은 문제에서 특정한 시작점을 지정하지 않았다는 점
  • 여러 시작점에서의 연결된 정점까지의 경로를 고려해야 함
• 모든 경로를 고려해 벨만-포드로 음의 사이클을 도출하도록 유니온-파인드 응용
  • 간선 정보를 입력받을 때 유니온-파인드 연산을 통해 연결된 두 정점을 한 집합으로 설정
  • 이후, 루트에 해당하는 정점들을 시작점으로 벨만-포드를 적용하고 음의 사이클이 존재하는 경우 true를 반환토록 로직 구성
  • 여러 시작점에 대한 벨만-포드 수행 결과 중 음의 사이클이 존재하는 경우가 하나라도 있는 경우 YES에 해당
• 로직의 시간복잡도는 평균적으로 $T\times N \times (M+W) = O(TN(M+W))$로 수렴. 주어진 제한 조건 2초를 무난히 통과

코드

import static java.lang.Integer.parseInt;

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

public class Main {

    static int[] parent;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = parseInt(br.readLine());
        List<String> answers = new ArrayList<>();

        while (T-- > 0) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int N = parseInt(st.nextToken());
            int M = parseInt(st.nextToken());
            int W = parseInt(st.nextToken());

            parent = new int[N + 1];
            Arrays.fill(parent, -1);

            List<Edge> edges = new ArrayList<>();
            while (M-- > 0) {
                st = new StringTokenizer(br.readLine());
                int u = parseInt(st.nextToken());
                int v = parseInt(st.nextToken());
                int w = parseInt(st.nextToken());

                edges.add(new Edge(u, v, w));
                edges.add(new Edge(v, u, w));
                union(u, v);
            }

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

                edges.add(new Edge(u, v, -w));
                union(u, v);
            }

            List<Integer> parents = new ArrayList<>();
            for (int i = 1; i < parent.length; i++) {
                if (parent[i] < 0) {
                    parents.add(i);
                }
            }
            List<Boolean> results = parents.stream().map(start -> isCycleExists(start, N, edges))
                .filter(result -> result)
                .collect(Collectors.toList());
            answers.add(results.isEmpty() ? "NO" : "YES");
        }

        System.out.println(answers.stream().collect(Collectors.joining("\n")));
        br.close();
    }

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

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

    static void union(int u, int v) {
        int p1 = find(u);
        int p2 = find(v);

        if (p1 == p2) {
            return;
        }

        if (parent[p1] < parent[p2]) {
            parent[p1] += parent[p2];
            parent[p2] = p1;
        } else {
            parent[p2] += parent[p1];
            parent[p1] = p2;
        }
    }

    static boolean isCycleExists(int start, int N, List<Edge> edges) {
        int[] dist = new int[N + 1];
        Arrays.fill(dist, Integer.MAX_VALUE);
        dist[start] = 0;

        for (int i = 0; i < N - 1; i++) {
            for (Edge e : edges) {
                if (dist[e.u] != Integer.MAX_VALUE && dist[e.v] > dist[e.u] + e.w) {
                    dist[e.v] = dist[e.u] + e.w;
                }
            }
        }

        for (Edge e : edges) {
            if (dist[e.u] != Integer.MAX_VALUE && dist[e.v] > dist[e.u] + e.w) {
                return true;
            }
        }

        return false;
    }

    static class Edge {

        int u, v, w;

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

결과