Union find

Leetcode 1562. Find Latest Group of Size M

Problem

Solution

We can use union find to union the neighbor 1s together into a big group. At the same time, we should maintain the size of each group.

typedef vector<int> vi;
class UF {
public:
vi p, setSize;
UF(int n) {
p.assign(n, 0); for (int i = 0; i < n; ++i) p[i] = i;
setSize.assign(n, 1);
}
int getSetSize(int i) {return setSize[i];}
int findSet(int i) {return p[i] == i ? i : (p[i] = findSet(p[i]));}
void unionSet(int i, int j) {
if (findSet(i) == findSet(j)) return;
int x = findSet(i), y = findSet(j);
p[x] = y;
setSize[y] += setSize[x];
return;
}
};
class Solution {
public:
int findLatestStep(vector<int>& arr, int m) {
const int N = arr.size();
vi nums(N, 0);
unordered_map<int, int> map_;
UF uf(N);
int ret = -1;
for (int i = 0; i < N; ++i) {
auto j = arr[i] - 1;
nums[j] = 1;
auto l = j-1, r = j+1;
if (l>=0 && nums[l] == 1) {
map_[uf.getSetSize(uf.findSet(l))] --;
uf.unionSet(l, j);
}
if (r < N && nums[r] == 1) {
map_[uf.getSetSize(uf.findSet(r))] --;
uf.unionSet(r, j);
}
map_[uf.getSetSize(uf.findSet(j))] ++;
if (map_[m]>0) ret = i+1;
}
return ret;
}
};

Data Scientist/MLE/SWE @takemobi