Union find

Leetcode 1562. Find Latest Group of Size M

Problem

1562. Find Latest Group of Size M

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;
    }
};