927-三等分

题目

927. 三等分

难度困难 51

给定一个由 0 和 1 组成的数组 A，将数组分成 3 个非空的部分，使得所有这些部分表示相同的二进制值。

如果可以做到，请返回任何 [i, j]，其中 i+1 < j，这样一来：

• A[0], A[1], ..., A[i] 组成第一部分；
• A[i+1], A[i+2], ..., A[j-1] 作为第二部分；
• A[j], A[j+1], ..., A[A.length - 1] 是第三部分。
• 这三个部分所表示的二进制值相等。

如果无法做到，就返回 [-1, -1]。

注意，在考虑每个部分所表示的二进制时，应当将其看作一个整体。例如，[1,1,0] 表示十进制中的 6，而不会是 3。此外，前导零也是被允许的，所以 [0,1,1] 和 [1,1] 表示相同的值。

示例 1：

输入：[1,0,1,0,1]
输出：[0,3]

示例 2：

输出：[1,1,0,1,1]
输出：[-1,-1]

提示：

1. 3 <= A.length <= 30000
2. A[i] == 0 或 A[i] == 1

思路

每组1的数目肯定相等。于是先统计1的位置并保存为bit，之后看bit的长度是否大于3、是否能被3整除。

能被3整除的情况下，分成三分，每份为t。每一份中1肯定都是t个，前缀0可以忽略，因此应该看后缀0有多少。前两份的前缀0和后面的后缀0混合在一起，可以自由组合，因此先不看。只看第3份的后缀0。如果第3份的后缀0个数比前两份之间的都多，也是不行的。在第三分的后缀0个数小于1、2和2、3之间的0的个数的时候，分别从第1份最后一个1后面延伸相应个数的0、第2份最后一个1后面延伸相应个数的0，返回位置即为答案。

代码

class Solution {
public:
    vector<int> threeEqualParts(vector<int>& arr) {
        vector<int> bit;
        int len = arr.size();
        for(int i=0;i<arr.size();i++)
        if(arr[i]==1)
        bit.push_back(i);
        if(bit.size()%3) return {-1,-1};
        if(bit.empty()) return {0,2};
        int t = bit.size()/3;
        int back_zero = len - 1 - bit.back();
        if(back_zero>bit[2*t]-bit[2*t-1]-1||back_zero>bit[t]-bit[t-1]-1)
        return {-1,-1};
        vector<int> first(bit.begin(),bit.begin()+t);
        vector<int> second(bit.begin()+t,bit.begin()+2*t);
        vector<int> third(bit.begin()+2*t,bit.end());
        for(int i=t-1;i>0;--i){
            if(first[i]-first[i-1]!=second[i]-second[i-1]||
            first[i]-first[i-1]!=third[i]-third[i-1]) 
            return {-1,-1};
        }
        int i = bit[t-1]+back_zero;
        int j = bit[2*t-1]+back_zero+1;
        return {i,j};
    }
};