LeetCode 803. Bricks Falling When Hit

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
class UnionFind {
private:
vector<int> f, sz;
public:
UnionFind(int n): f(n), sz(n) {
for (int i = 0; i < n; ++i) {
f[i] = i;
sz[i] = 1;
}
}
int find(int x) {
return f[x] == x ? x : f[x] = find(f[x]);
}
void merge(int x, int y) {
int fx = find(x), fy = find(y);
if (fx != fy) {
if (sz[fx] > sz[fy])
swap(fx, fy);
sz[fy] += sz[fx];
f[fx] = fy;
}
}
int size(int x) {
return sz[find(x)];
}
};
class Solution {
public:
vector<int> hitBricks(vector<vector<int>>& grid, vector<vector<int>>& hits) {
const int R = grid.size(), C = grid[0].size();
UnionFind uf(R * C + 1);
vector<vector<int>> status = grid;
for (int i = 0; i < hits.size(); ++i)
status[hits[i][0]][hits[i][1]] = 0;
for (int r = 0; r < R; ++r)
for (int c = 0; c < C; ++c)
if (status[r][c] == 1) {
if (r == 0)
uf.merge(R * C, r * C + c);
if (r > 0 && status[r - 1][c] == 1)
uf.merge(r * C + c, (r - 1) * C + c);
if (c > 0 && status[r][c - 1] == 1)
uf.merge(r * C + c, r * C + c - 1);
}
const vector<int> dirs{-1, 0, 1, 0, -1};
vector<int> ans(hits.size(), 0);
for (int i = hits.size() - 1; i >= 0; --i) {
int r = hits[i][0], c = hits[i][1];
if (grid[r][c] == 0) continue;
int prev = uf.size(R * C);
if (r == 0)
uf.merge(c, R * C);
for (int j = 0; j < 4; ++j) {
int nr = r + dirs[j], nc = c + dirs[j + 1];
if (nr < 0 || nr >= R || nc < 0 || nc>= C || status[nr][nc] != 1) continue;
uf.merge(r * C + c, nr * C + nc);
}
int curr = uf.size(R * C);
ans[i] = max(0, curr - prev - 1);
status[r][c] = 1;
}
return ans;
}
};