Find Latest Group of Size M

Given an array arr that represents a permutation of numbers from 1 to n. You have a binary string of size n that initially has all its bits set to zero.

At each step i (assuming both the binary string and arr are 1-indexed) from 1 to n, the bit at position arr[i] is set to 1. You are given an integer m and you need to find the latest step at which there exists a group of ones of length m. A group of ones is a contiguous substring of 1s such that it cannot be extended in either direction.

Return the latest step at which there exists a group of ones of length exactly m. If no such group exists, return -1.

 

Example 1:

Input: arr = [3,5,1,2,4], m = 1
Output: 4
Explanation:
Step 1: "00100", groups: ["1"]
Step 2: "00101", groups: ["1", "1"]
Step 3: "10101", groups: ["1", "1", "1"]
Step 4: "11101", groups: ["111", "1"]
Step 5: "11111", groups: ["11111"]
The latest step at which there exists a group of size 1 is step 4.

思路：Union Find,我觉得这题最难的是想起用count数组去记录size的个数状态，最后直接返回count[m] >= 1 状态；

class Solution {
    private class UnionFind {
        public int[] father;
        public int[] size;
        public int[] count;
        public int m;
        public UnionFind (int n) {
            this.father = new int[n + 1];
            this.size = new int[n + 1];
            this.count = new int[n + 1];
            for(int i = 0; i < n; i++) {
                father[i] = i;
            }
        }    
        
        public int find(int x) {
            int j = x;
            while(father[j] != j) {
                j = father[j];
            }
            
            // path compression;
            while(x != j) {
                int fx = father[x];
                father[x] = j;
                x = fx;
            }
            return j;
        }
        
        public void union(int a, int b) {
            int root_a = find(a);
            int root_b = find(b);
            if(root_a != root_b) {
                father[root_a] = root_b;
                --count[size[root_a]];
                --count[size[root_b]];
                
                size[root_b] += size[root_a];
                size[root_a] = 0;
                
                ++count[size[root_b]];
            }
        }
        
        public int getCount(int m) {
            return count[m];
        }
        
        public void setSize(int a) {
            size[a] = 1;
            ++count[1];
        }
    }
    
    public int findLatestStep(int[] arr, int m) {
        int laststep = -1;
        int n = arr.length;
        UnionFind uf = new UnionFind(n);
        int[] array = new int[n];
        for(int i = 0; i < n; i++) {
            int index = arr[i];
            array[index - 1] = 1;
            uf.setSize(index - 1);
            if(index < n && array[index] == 1) {
                uf.union(index - 1, index);
            }
            if(index - 2 >= 0 && array[index - 2] == 1) {
                uf.union(index - 1, index - 2);
            }
            if(uf.getCount(m) >= 1) {
                laststep = i + 1;
            }
        }
        return laststep;
    }
}