012.md

012 - Red Painting (★4)

解答

#include <iostream>
#include <vector>
#include <numeric> // std::iota()

// Union-Find 木
class UnionFind
{
public:

	UnionFind() = default;

	// n 個の要素
	explicit UnionFind(size_t n)
		: m_parents(n)
	{
		std::iota(m_parents.begin(), m_parents.end(), 0);
	}

	// i の root を返す
	int find(int i)
	{
		if (m_parents[i] == i)
		{
			return i;
		}

		// 経路圧縮
		return (m_parents[i] = find(m_parents[i]));
	}

	// a の木と b の木を統合
	void merge(int a, int b)
	{
		a = find(a);
		b = find(b);

		if (a != b)
		{
			m_parents[b] = a;
		}
	}

	// a と b が同じ木に属すかを返す
	bool connected(int a, int b)
	{
		return (find(a) == find(b));
	}

private:

	// m_parents[i] は i の 親,
	// root の場合は自身が親
	std::vector<int> m_parents;
};

// (r, c) を一次元配列のインデックスに変換.
// 範囲外の場合は -1 を返す
int ToIndex(int r, int c, int H, int W)
{
	if ((r < 0) || (H <= r) || (c < 0) || (W <= c))
	{
		return -1;
	}

	return (W * r + c);
}

int main()
{
	// H 行 W 列
	int H, W;
	std::cin >> H >> W;

	// Q 個のクエリ
	int Q;
	std::cin >> Q;

	// Union Find 木の作成
	UnionFind uf(H * W);

	// 赤で塗られているか, 初期値はすべて false
	std::vector<bool> red(H * W);

	// 各クエリについて
	for (int q = 0; q < Q; ++q)
	{
		// クエリの種類
		int t;
		std::cin >> t;

		if (t == 1)
		{
			int r, c;
			std::cin >> r >> c;
			--r; --c;

			const int from = ToIndex(r, c, H, W);
			constexpr int dx[4] = { 0, -1, 1, 0 };
			constexpr int dy[4] = { -1, 0, 0, 1 };

			// 赤色で塗る
			red[from] = true;

			// 上, 左, 右, 下のマスについて
			for (int i = 0; i < 4; ++i)
			{
				const int to = ToIndex(r + dy[i], c + dx[i], H, W);

				// そのマスが赤で塗られていたら, そのマスと
				if (to != -1 && red[to])
				{
					uf.merge(from, to);
				}
			}
		}
		else
		{
			int ra, ca, rb, cb;
			std::cin >> ra >> ca >> rb >> cb;
			--ra; --ca; --rb; --cb;

			const int a = ToIndex(ra, ca, H, W);
			const int b = ToIndex(rb, cb, H, W);

			// 二か所とも赤で塗られていて, 同じ root を持つ場合
			if (red[a] && red[b] && uf.connected(a, b))
			{
				std::cout << "Yes\n";
			}
			else
			{
				std::cout << "No\n";
			}
		}
	}
}