leetcode-go/medium/200.go

package main

// Grid coordinate of single grid
type Grid struct {
	X int
	Y int
}

// UF state of union-find algorithm
type UF struct {
	ID   map[Grid]Grid
	Size map[Grid]int
	Cnt  int
}

// Union ...
func (uf *UF) Union(x1, y1, x2, y2 int) {
	g1 := uf.Find(x1, y1)
	g2 := uf.Find(x2, y2)
	if g1 == g2 {
		return
	}
	if uf.Size[g1] < uf.Size[g2] {
		uf.ID[g1] = g2
		uf.Size[g2] += uf.Size[g1]
	} else {
		uf.ID[g2] = g1
		uf.Size[g1] += uf.Size[g2]
	}
	uf.Cnt--
}

// Find ...
func (uf *UF) Find(x, y int) Grid {
	g := Grid{x, y}
	for g != uf.ID[g] {
		g = uf.ID[g]
	}
	return g
}

// Connected ...
func (uf *UF) Connected(x1, y1, x2, y2 int) bool {
	return uf.Find(x1, y1) == uf.Find(x2, y2)
}

func numIslands(grid [][]byte) int { // Union-find algorithm
	hei := len(grid)
	if hei == 0 {
		return 0
	}
	wid := len(grid[0])
	var uf UF
	uf.ID = make(map[Grid]Grid)
	uf.Size = make(map[Grid]int)
	for y := 0; y < hei; y++ {
		for x := 0; x < wid; x++ {
			if grid[y][x] == '1' {
				g := Grid{x, y}
				uf.ID[g] = g
				uf.Size[g] = 1
				uf.Cnt++
			}
		}
	}
	for y := 0; y < hei; y++ {
		for x := 0; x < wid; x++ {
			if grid[y][x] == '0' {
				continue
			}
			if x+1 < wid && grid[y][x+1] == '1' {
				uf.Union(x, y, x+1, y)
			}
			if y+1 < hei && grid[y+1][x] == '1' {
				uf.Union(x, y, x, y+1)
			}
		}
	}
	return uf.Cnt
}

func main() {
	println(numIslands([][]byte{
		{'1', '1', '0', '0', '0'},
		{'1', '1', '0', '0', '0'},
		{'0', '0', '1', '0', '0'},
		{'0', '0', '0', '1', '1'}}))
}