012 - Red Painting（★4）

キーワード

解説

UnionFind を使うことで, $\mathrm{O}(HW + Q\alpha(HW))$ で解くことができる.

色を塗るたびに, 左右上下を確認して,すでに塗られていればuniteしていけば良い.

提出コード

func main() {
	var h, w int
	scanIntVariables(&h, &w)
	Q := readInts()[0]
 
	g := make([][]int, h+1)
	for i := 0; i <= h; i++ {
		g[i] = make([]int, w+1)
	}
 
	ids := make([][]int, h+1)
	for i := 0; i <= h; i++ {
		ids[i] = make([]int, w+1)
	}
	n := 0
	for i := 0; i <= h; i++ {
		for j := 0; j <= w; j++ {
			ids[i][j] = n
			n++
		}
	}
 
	uf := newUnionFind((h + 1) * (w + 1))
	dx := []int{0, 1, 0, -1}
	dy := []int{-1, 0, 1, 0}
	for i := 0; i < Q; i++ {
		q := readInts()
		if q[0] == 1 {
			r, c := q[1], q[2]
			g[r][c] = 1
			for j := 0; j < 4; j++ {
				nr, nc := r+dy[j], c+dx[j]
				if nr <= 0 || nr > h || nc <= 0 || nc > w {
					continue
				}
				if g[nr][nc] == 1 {
					uf.unite(ids[nr][nc], ids[r][c])
				}
			}
		} else {
			r1, c1, r2, c2 := q[1], q[2], q[3], q[4]
			if g[r1][c1] == 0 || g[r2][c2] == 0 {
				writeLine("No")
				continue
			}
			if uf.isSame(ids[r1][c1], ids[r2][c2]) {
				writeLine("Yes")
			} else {
				writeLine("No")
			}
		}
	}
}
 
type UnionFind struct {
	n    int
	root []int
}
 
func newUnionFind(n int) UnionFind {
	root := make([]int, n)
	for i := 0; i < n; i++ {
		root[i] = -1
	}
	uf := UnionFind{n: n, root: root}
	return uf
}
func (u *UnionFind) find(x int) int {
	if u.root[x] < 0 {
		return x
	}
	u.root[x] = u.find(u.root[x])
	return u.root[x]
}
func (u *UnionFind) unite(x, y int) {
	x = u.find(x)
	y = u.find(y)
	if x == y {
		return
	}
	if -u.root[x] < -u.root[y] {
		x, y = y, x
	} // xの方がサイズ大きいように
	u.root[x] += u.root[y]
	u.root[y] = x
}
func (u *UnionFind) isSame(x, y int) bool {
	return u.find(x) == u.find(y)
}
func (u *UnionFind) size(x int) int {
	return -u.root[u.find(x)]
}

典型90 022 - Cubic Cake（★2）