[Java] 백준 13418번 : 학교 탐방하기

1. 문제

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

13418번: 학교 탐방하기

 

 

2. 풀이과정

MST 알고리즘은 kruskal, prim이 있다. 여기서는 간선의 수가 적으므로 kruskal을 사용했다.

 

처음에는 전체 edge를 입력으로 받아서 오름차순 정렬 후, 최대 & 최소 MST를 찾았다.

하지만, 이 문제에서는 cost가 0, 1 두 가지로 주어지기 때문에 정렬을 하지 않고 오르막길과 내리막길을 따로 모아서 최대로 만들 수 있는 MST를 구하는 방식으로 풀면 훨씬 효율적이다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
import java.util.ArrayList;

public class Main {

    static class Edge implements Comparable<Edge> {
        int n1, n2, cost;

        public Edge(int n1, int n2, int cost){
            this.n1 = n1;
            this.n2 = n2;
            this.cost = cost;
        }

        @Override
        public int compareTo(Edge o){
            return this.cost - o.cost;
        }
    }
    static int N, p[];

    public static int find(int n1){
        if(p[n1] == n1){
            return n1;
        }
        return p[n1] = find(p[n1]);
    }

    public static int mst(ArrayList<Edge> edges){
        int cnt = 0;
        p = new int[N+1];
        for(int i=1; i<=N; i++){p[i] = i;}

        for(Edge edge : edges){
            int n1 = find(edge.n1);
            int n2 = find(edge.n2);
            if(n1==n2){ continue;}
            p[n2] = n1;
            if(++cnt == N){break;}
        }

        return cnt;
    }

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken())+1;
        ArrayList<Edge> up = new ArrayList<>();
        ArrayList<Edge> down = new ArrayList<>();

        for(int i=0; i<M; i++){
            st = new StringTokenizer(br.readLine());
            int n1 = Integer.parseInt(st.nextToken());
            int n2 = Integer.parseInt(st.nextToken());
            int cost = Integer.parseInt(st.nextToken());
            if(cost==0){ up.add(new Edge(n1, n2, cost));}
            else{ down.add(new Edge(n1, n2, cost));}
        }

        int min = N-mst(down);
        int max = mst(up);
        System.out.println(max*max - min*min);
    }
}

 

결과