并查集

Posted on Edited on In Word count in article: 7.7k Reading time ≈ 7 mins.
go的并查集实现，以及相关问题

并查集

代码模板

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
type UF struct {
	union []int
}

// 初始化
func newUF(cap int) *UF {
	uf := UF{
		make([]int, cap),
	}
	// 每个元素的索引等于自己,每一个元素是一个集合，索引也代表元素
	// 每个元素是代表元素
	for i := 0; i < cap; i++ {
		uf.union[i] = i
	}
	return &uf
}

// 合并
func (u *UF) Union(x, y int) {
	// 查找两个元素的代表元素
	rootX := u.find(x)
	rootY := u.find(y)
	// 如果代表元素一样 ，则属于同一集合
	if rootX == rootY {
		return
	}
	// 如果不是则x元素的代表元素位置的代表元素改为y的代表元素
	u.union[rootX] = rootY
}
// 判断两个元素是不是同一集合。也就是两个位置上的root代表元素是不是同一个
func (u *UF) Connected(x, y int) bool {
	return u.find(x) == u.find(y)
}

// 查询一个元素的代表元素，属于哪个集合
func (u *UF) find(x int) int {
	root := x
	// x元素位置是不是x自己，如果不是，说明该元素被合并过，有别的代表元素
	// 被合并过后x位置上的元素就是它的父节点
	// 此过程是不断往上找 父节点，找他的代表元素
	for u.union[root] != root {
		root = u.union[root]
	}
	// 压缩路径，将该条分支上的所有子节点的位置上的元素都是root
	// 先取x位置上的节点，然后x位置上改为root
	// 再把x位置改为刚取出的节点的位置，看是不是root
	for x != root {
		tmp := u.union[x]
		u.union[x] = root
		x = tmp
	}
	// 最后返回root节点，也就是x的代表元素
	return root
}

岛屿数量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
// dfs
func numIslands(grid [][]byte) int {
    var dfs func(i,j int)
    dfs = func(i, j int) {
        if i >= 0 && j >= 0 && i < len(grid) && j < len(grid[0]) && grid[i][j] == '1' {
            grid[i][j] = '0'
            dfs(i + 1, j)
            dfs(i - 1, j)
            dfs(i, j - 1)
            dfs(i, j + 1)
        }
    }
    
    res := 0
    for i := 0; i < len(grid); i++ {
        for j := 0; j < len(grid[0]); j++ {
            if grid[i][j] == '1' {
                dfs(i, j)
                res++
            }
        }
    }
    return res
}

// 并查集方法
// 并查集，记录count
type UF struct {
	union []int
    count int
}

// 初始化并查集
func newUF(m, n int, grid [][]byte) *UF {
	uf := UF{make([]int, m * n), 0}
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
        	// 默认所有的1都是单独集合，索引位置是i*n+j
            uf.union[i * n + j] = i * n + j
            // 如果是1，集合数+1
            if grid[i][j] == '1' {
                uf.count++
            }
        }
    }
	return &uf
}
// 合并，合并成功，集合数-1
func (u *UF) Union(x, y int) {
	rootX := u.find(x)
	rootY := u.find(y)
	if rootX == rootY {
		return
	}
	u.union[rootX] = rootY
    u.count--
}
//  查找不需要修改
func (u *UF) find(x int) int {
	root := x
	for u.union[root] != root {
		root = u.union[root]
	}
	for x != root {
		tmp := u.union[x]
		u.union[x] = root
		x = tmp
	}
	return root
}
// 主函数
func numIslands(grid [][]byte) int {
    m := len(grid)
    n := len(grid[0])
    uf := newUF(m, n, grid)

    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
        	// 如果是1，则看其他四个方向是不是1，是1就合并，然后置成0
            if grid[i][j] == '1' {
                if i - 1 >= 0 && grid[i-1][j] == '1' {
					uf.Union(i*n+j, (i-1)*n+j)
				}
				if i + 1 < m && grid[i+1][j] == '1' {
					uf.Union(i*n+j, (i+1)*n+j)
				}
				if j - 1 >= 0 && grid[i][j-1] == '1' {
					uf.Union(i*n+j, i*n+(j-1))
				}
				if j + 1 < n && grid[i][j+1] == '1' {
					uf.Union(i*n+j, i*n+(j+1))
				}
				grid[i][j] = '0'
            }
        }
    }
    // 最后返回集合数即可
    return uf.count
}

被围绕区域

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
// dfs
func solve(board [][]byte)  {
    var dfs func(i, j int)
    dfs = func(i, j int) {
        if i >= 0 && j >= 0 && i < len(board) && j < len(board[0]) {
            if board[i][j] == 'O' {
            	// 先改成别的，最后再遍历统一更新
                board[i][j] = 'Y'
                dfs(i - 1, j)
                dfs(i + 1, j)
                dfs(i, j - 1)
                dfs(i, j + 1)
            } 
        } 
    }

    for i := 0; i < len(board); i++ {
        if board[i][0] == 'O' { dfs(i, 0) }
        if board[i][len(board[0]) - 1] == 'O' { dfs(i, len(board[0]) - 1) }
    }
    for j := 0; j < len(board[0]); j++ {
        if board[0][j] == 'O' { dfs(0, j) }
        if board[len(board) - 1][j] == 'O' { dfs(len(board) - 1, j) }
    }
    for i := 0; i < len(board); i++ {
        for j := 0; j < len(board[0]); j++ {
            if board[i][j] != 'Y' {
                board[i][j] = 'X'
            } else {
                board[i][j] = 'O'
            }
        }
    }
}

// 并查集
type UF struct {
    parents []int
}

func newUF (n int) *UF {
    uf := UF{make([]int, n)}
    for i := 0; i < n; i++ {
        uf.parents[i] = i
    }
    return &uf
}

func(uf *UF)find(x int) int {
    root := x
    for root != uf.parents[root] {
        root = uf.parents[root]
    }
    for x != root {
        tmp := uf.parents[x]
        uf.parents[x] = root
        x = tmp
    }
    return root
}

func(uf *UF)union(x, y int) {
    rootX := uf.find(x)
    rootY := uf.find(y)
    if rootX == rootY {
        return
    }
    uf.parents[rootX] = rootY
}

func(uf *UF)isConnect(x, y int) bool {
    return uf.find(x) == uf.find(y)
}

func solve(board [][]byte)  {
    m := len(board)
    n := len(board[0])
    t := m * n
    // 多创建一个集合，边界元素放到该集合内，以该集合为准
    uf := newUF(t + 1)
  
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            if board[i][j] == 'O' {
                // 边界位置都归到一个单独的集合
                if i == 0 || j == 0 || i == m - 1|| j == n -1 {
                    uf.union(i * n + j, t)
                }
                // 只判断左边和上边，即可覆盖全部元素
                if i > 0 && board[i - 1][j] == 'O' { uf.union(i *n + j, (i - 1)*n + j) }
                if j > 0 && board[i][j - 1] == 'O' { uf.union(i *n + j, i*n + j - 1) }
            }
        }
    }
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            // 判断是O的如果不和边界的元素属于同一集合，，则是被包围，更新为X
            if board[i][j] == 'O' && !uf.isConnect(i * n + j, t) {
                board[i][j] = 'X'
            }
        }
    }
}

省份数量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
// dfs
func findCircleNum(isConnected [][]int) int {
    var dfs func(i, j int)
    dfs = func(i,j int) {
        if isConnected[i][j] == 1 {
            isConnected[i][j] = 0
            isConnected[j][i] = 0
            for x := 0; x < len(isConnected); x++ {
                dfs(x, j)
            }
            for y := 0; y < len(isConnected[0]); y++ {
                dfs(i, y)
            }
        }
    }

    res := 0
    for i := 0; i < len(isConnected); i++ {
        for j := 0; j < len(isConnected[0]); j++ {
            if isConnected[i][j] == 1 {
                dfs(i, j)
                res++
            }
        }
    }
    return res
}

// 并查集1，每个元素单独处理，所有元素都加入集合
type UF struct {
    union []int
    count int
}

func newUf (m, n int, isConnected [][]int) *UF {
    uf := UF{make([]int, m * n), 0}
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            uf.union[i * n + j] = i * n + j
            if isConnected[i][j] == 1 {
                uf.count++
            }
        }
    }
    return &uf
}

func(uf *UF) find(x int) int {
    root := x
    for uf.union[root] != root {
        root = uf.union[root]
    }
    for x != root {
        tmp := uf.union[x]
        uf.union[x] = root
        x = tmp
    }
    return root
}

func(uf *UF) Union(x, y int) {
    rootX := uf.find(x)
    rootY := uf.find(y)
    if rootX == rootY {
        return
    }
    uf.union[rootX] = rootY
    uf.count--
}

func findCircleNum(isConnected [][]int) int {
    m := len(isConnected)
    n := len(isConnected[0])
    uf := newUf(m, n, isConnected)
    // 看全部位置
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            if isConnected[i][j] == 1 {
            	// 合并对称位置
                uf.Union(i*n + j, j*n + i)
                isConnected[i][j] = 0
                // 合并当前位置和对称位置的同一列
                for x := 0; x < m; x++ {
                    if isConnected[x][j] == 1 {
                        uf.Union(i*n + j, x*n + j)
                    }
                    if isConnected[x][i] == 1{
                        uf.Union(j*n + i, x*n + i)
                    }
                }
                // 合并当前位置和对称位置的同一行
                for y := 0; y < n; y++ {
                    if isConnected[i][y] == 1 {
                        uf.Union(i*n + j, i*n + y)
                    }
                    if isConnected[j][y] == 1 {
                        uf.Union(j*n + i, j*n + y)
                    }
                }
            }
        }
    }
    return uf.count
}

// 并查集2， 优化，初始化并查集直接按对称位置初始化
type UF struct {
    union []int
    count int
}

func newUf (m int) *UF {
    uf := UF{make([]int, m), m}
    for i := 0; i < m; i++ {
        uf.union[i] = i
    }
    return &uf
}

func(uf *UF) Union(x, y int) {
    rootX := uf.find(x)
    rootY := uf.find(y)
    if rootX == rootY {
        return
    }
    uf.union[rootX] = rootY
    uf.count--
}


func(uf *UF) find(x int) int {
    root := x
    for uf.union[root] != root {
        root = uf.union[root]
    }
    for x != root {
        tmp := uf.union[x]
        uf.union[x] = root
        x = tmp
    }
    return root
}

func findCircleNum(isConnected [][]int) int {
    m := len(isConnected)
    uf := newUf(m)
    // 只看对称位置
    for i := 0; i < m; i++ {
        for j := i; j < m; j++ {
            if isConnected[i][j] == 1 {
                uf.Union(i, j)
            }
        }
    }
    return uf.count
}