Codeforces Round #309 Div1 C: Love Triangles

問題

codeforces.com

解法

mayokoex.hatenablog.com

問題分が分かりにくいが、「ある2人が同じ人を好き、あるいは嫌いな場合、その2人は愛しあわなければならない」を「グループ内の人間関係がloveのみになるようにグループ分け出来る」と理解する。

まずはUnion-Findを使ってlove同士を結合し、hateなのに結合されている2点がないかどうか調べる。あった場合は題意を満たさないので0通りになる。

次に、結合された点の集合(グループ)を1つの頂点と見なして、グループ間の関係を表すグラフを作る。この時、love同士は全て結合されているので、グループ間の関係が存在する場合、それはhateのみである。

グループ間のグラフができたら、2色で塗り分けて矛盾がないかどうか調べる(二部グラフかどうか調べる)。例えばAグループとBグループが互いにhateで、BグループとCグループが互いにhateなら、AグループとCグループは互いにloveでなければならない。

ここまで矛盾が無かったら、個人ではなくグループについてUnion-Find木を作り、グループをグループ化する。上記のようにAとBが互いにhateで、BとCが互いにhateなら、AとCの関係はLoveに固定されるので、この3グループの関係性は1通りである。結合し終わると、大きなグループ(グループ化されたグループ)の間には関係は存在せず、LoveでもHateでもあり得る状態になる。よってこれらの関係の組み合わせは、大きなグループの数をxとした時、 2^{(x-1)}通りになる。

コード

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Scanner;
import java.util.TreeSet;

public class Main {

	private final int MOD = 1000000007;
	private int[] color;// 二部グラフ化どうか判定するために各ノードを2色で塗り分ける
	private ArrayList<TreeSet<Integer>> hateGraph;

	public void solve() {
		Scanner sc = new Scanner(System.in);
		int N = sc.nextInt();
		int M = sc.nextInt();
		color = new int[N];

		UnionFind uf = new UnionFind(N);
		int[] a = new int[M];
		int[] b = new int[M];
		int[] c = new int[M];

		for (int i = 0; i < M; i++) {
			a[i] = sc.nextInt() - 1;
			b[i] = sc.nextInt() - 1;
			c[i] = sc.nextInt();

			// loveなら同じグループにしておく
			if (c[i] == 1) {
				uf.unite(a[i], b[i]);
			}
		}

		for (int i = 0; i < M; i++) {
			// hateなのに同じグループにいる時
			if (c[i] == 0 && uf.isSame(a[i], b[i])) {
				System.out.println(0);
				return;
			}
		}

		// union-find内で記録したグループ番号を連番に直しておく
		HashMap<Integer, Integer> groupIdMap = new HashMap<>();
		int groupNum = 0;
		for (int i = 0; i < N; i++) {
			if (!groupIdMap.containsKey(uf.find(i))) {
				groupIdMap.put(uf.find(i), groupNum);
				groupNum++;
			}
		}

		// hate関係のみで構成されたグループ間の関係を示すグラフを作る
		hateGraph = new ArrayList<>();
		for (int i = 0; i < groupNum; i++) {
			hateGraph.add(new TreeSet<Integer>());
		}
		for (int i = 0; i < M; i++) {
			if (c[i] == 0) {
				int aGroupId = groupIdMap.get(uf.find(a[i]));
				int bGroupId = groupIdMap.get(uf.find(b[i]));
				hateGraph.get(aGroupId).add(bGroupId);
				hateGraph.get(bGroupId).add(aGroupId);
			}
		}

		// 二部グラフになっているかどうか
		for (int i = 0; i < groupNum; i++) {
			if (color[i] == 0) {
				if (!colorDFS(i, 1)) {
					System.out.println(0);
					return;
				}
			}
		}

		// グループ間の関係について固定されているものを結合する
		UnionFind groupFind = new UnionFind(groupNum);
		for (int i = 0; i < M; i++) {
			if (c[i] == 0) {
				int aGroupId = groupIdMap.get(uf.find(a[i]));
				int bGroupId = groupIdMap.get(uf.find(b[i]));
				groupFind.unite(aGroupId, bGroupId);
			}
		}

		long ans = 1;
		for (int i = 0; i < groupFind.size() - 1; i++) {
			ans *= 2;
			ans %= MOD;
		}
		System.out.println(ans);

	}

	private boolean colorDFS(int groupId, int colorId) {
		color[groupId] = colorId;
		// groupIdをhateなグループについて、-colorIdで塗る
		for (int hate : hateGraph.get(groupId)) {
			if (color[hate] == colorId) {
				// hateなのに同じ色にいたらダメ
				return false;
			}
			if (color[hate] == -colorId) {
				// 既に塗り終わってたらok
				continue;
			}
			if (!colorDFS(hate, -colorId)) {
				// まだ塗ってなかったら、こいつについてDFSして、そこでダメだったら終了
				return false;
			}
		}
		// 問題なければ正常終了
		return true;

	}

	public static void main(String[] args) {
		new Main().solve();
	}

}

class UnionFind {
	int[] parts;
	private int size;

	UnionFind(int n) {
		size = n;
		parts = new int[n];
		for (int i = 0; i < n; i++) {
			parts[i] = i;
		}
	}

	public int find(int x) {
		if (parts[x] == x)
			return x;
		return parts[x] = find(parts[x]);
	}

	public Boolean isSame(int x, int y) {
		return find(x) == find(y);
	}

	public void unite(int x, int y) {
		if (find(x) == find(y))
			return;
		size--;
		parts[find(x)] = find(y);
	}

	public int size() {
		return size;
	}
}