LeetCode打卡

• 927.三等分

927.三等分

https://leetcode.cn/problems/three-equal-parts/

class Solution {
public:
    vector<int> threeEqualParts(vector<int>& arr) {
        int total = 0, len = arr.size();
        for(int i = 0 ; i < len; i ++) {
            if(arr[i]) total ++;
        }

        if(total == 0) return {0, 2};
        if(total % 3 != 0) return {-1, -1};

        int avg = total / 3;
        int f[6] = {1, avg, avg + 1, 2*avg, 2*avg + 1, total};
        int site[6];

        int cnt1 = 0;
        for(int i = 0, j = 0 ; i < len; i ++) {
            if(!arr[i]) continue;
            cnt1 ++;
            while(j < 6 && cnt1 == f[j]) site[j ++] = i; //
        }

        int cnt0 = len - 1 - site[5];
        // 判断前两段尾部零是否够用
        if(site[2]-site[1]-1 < cnt0 || site[4]-site[3]-1 < cnt0) return {-1, -1};

        // 判断三段是否完全相同
        if(!judge(arr, site[0], site[1], site[2], site[3])) return {-1, -1};
        if(!judge(arr, site[0], site[1], site[4], site[5])) return {-1, -1};
        return {site[1]+cnt0, site[3]+cnt0+1};
    }

    bool judge(vector<int>& arr, int l1, int r1, int l2, int r2){
        for(int i = l1, j = l2; i <= r1; i ++, j ++)
            if(arr[i]^arr[j])
                return false;
        
        return true;
    }
};