MST 게임


백준 [16202] MST 게임

문제보기
Alt text

소스코드

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

public class Main {
	static int N, M, K;
	static int[][] edges;
	static int[] parents;
	static int[] rank;

	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		N = Integer.parseInt(st.nextToken());
		M = Integer.parseInt(st.nextToken());
		K = Integer.parseInt(st.nextToken());
		edges = new int[M][3];

		for (int i = 0; i < M; i++) {
			st = new StringTokenizer(br.readLine());
			edges[i][0] = Integer.parseInt(st.nextToken()) - 1;
			edges[i][1] = Integer.parseInt(st.nextToken()) - 1;
			edges[i][2] = i + 1;
		}

		Arrays.sort(edges, new Comparator<int[]>() {
			@Override
			public int compare(int[] o1, int[] o2) {
				// TODO Auto-generated method stub
				return Integer.compare(o1[2], o2[2]);
			}
		});

		chkMST(0);
	}

	public static void chkMST(int cnt) {
		if (cnt == K) {
			return;
		}
		parents = new int[N];
		rank = new int[N];
		for (int i = 0; i < N; i++) {
			parents[i] = i;
		}

		int sum = 0;
		int edgeCnt = 0;
		boolean flag = false;
		for (int i = 0; i < M; i++) {
			if (edges[i][2] == -1)
				continue;
			int px = findSet(edges[i][0]);
			int py = findSet(edges[i][1]);
			int cost = edges[i][2];

			if (px != py) {
				edgeCnt++;
				sum += cost;
				union(px, py);
			}
			if (edgeCnt == N - 1) {
				flag = true;
				break;
			}
		}

		if (!flag) {
			for (int i = cnt; i < K; i++) {
				System.out.print(0 + " ");
			}
			return;
		} else {
			edges[cnt][2] = -1;
			System.out.print(sum + " ");
			chkMST(cnt + 1);
		}
	}

	public static void makeSet(int x) {
		parents[x] = x;
	}

	public static int findSet(int x) {
		if (parents[x] == x)
			return x;
		parents[x] = findSet(parents[x]);
		return parents[x];
	}

	public static void union(int px, int py) {
		if (rank[px] > rank[py])
			parents[py] = px;
		else {
			parents[px] = py;
			if (rank[px] == rank[py])
				rank[py]++;
		}
	}
}