[백준 17398 JAVA] 통신망 분할

단순하게 본다면 dfs bfs로 풀리는 문제

 

처음 시도때는 그래프와 dfs로 구현했지만 당연하게도 시간초과가 뜨고말았다.

 

union find를 응용해서 풀어야 하는 문제이다.

 

제거하는 연결을 제외하고 미리 Union find를 통해서 전부 이어준다.

 

root에는 초기 값으로 전부 -1을 넣어준다 (음수 *음수 = 음수, 본인 노드 포함 1개)

 

union을 하면서 Root노드에 계속해서 갯수를 더해준다. (음수그대로 더한다)

Root노드가 아닌 노드에는 root노드 번호 가 들어가 있다.

 

제거되는 연결을 반대 순서로 결합해 준다.

이때 x,y가 Root노드가 같다면 비용이 0이고,

다르다면, root노드에 저장되어있는 하위 노드의 갯수를 곱해주면 된다.

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Stack;
import java.util.StringTokenizer;

class Main{
    static int stoi(String string){
        return Integer.parseInt(string);
    }
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = stoi(st.nextToken());
        int M = stoi(st.nextToken());
        int Q = stoi(st.nextToken());

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

        ArrayList<int[]> list = new ArrayList<>();
        list.add(new int[] {0,0});

        for(int i =0; i<M; i++){
            st = new StringTokenizer(br.readLine());
            list.add(new int[]{stoi(st.nextToken()),stoi(st.nextToken())});
        }

        boolean[] visited = new boolean[M+1];
        Stack<Integer> stack = new Stack<>();
        for(int i =0; i<Q; i++){
            int query = stoi(br.readLine());
            stack.add(query);
            visited[query] = true;
        }
        for(int i = 1; i<=M;i++){
            if(!visited[i]){
                int [] temp = list.get(i);
                union(temp[0],temp[1],root);
            }
        }
        long answer = 0;

        for(int i=0; i<Q; i++){
            int query = stack.pop();
            int[] input = list.get(query);

            int x = find(input[0],root);
            int y = find(input[1],root);
            if (x!=y){
                answer += (long) root[x] *root[y];
            }
            union(x,y,root);
        }
        System.out.println(answer);
    }
    public static void union(int a, int b, int[] root){
        a = find(a,root);
        b = find(b,root);
        if(a==b) return;

        root[a]+=root[b];
        root[b] = a;
    }

    public static int find(int a,  int[] root){
        if(root[a]<0) return a;
        return root[a] = find(root[a], root);
    }
}