Ba phần bằng nhau trong C++

Cho mảng A chỉ chứa 0 và 1. Chia thành 3 phần không rỗng có cùng giá trị nhị phân. Trả về [i, j] với i+1 < j. Nếu không thể trả về [-1, -1].

Ví dụ: [0,1,0,1,1] → [0,4] (các phần: [0], [1,0], [1,1] đều bằng 1 nhị phân).

Giải thuật

1. Đếm tổng số bit 1 (num). Nếu num % 3 != 0 → [-1,-1]. Nếu num == 0 → [0,2].
2. Mỗi phần cần req = num/3 bit 1. Tìm vị trí bắt đầu của phần 3 (idx) bằng cách đếm ngược req bit 1 từ cuối.
3. Hàm getIdx(left, right): bỏ qua số 0 đầu, sau đó so sánh bit một một từ left và right. Trả về left-1 nếu khớp hết, ngược lại -1.
4. firstEnd = getIdx(0, idx), secondEnd = getIdx(firstEnd+1, idx). Trả về {firstEnd, secondEnd+1}.

Ví dụ mã C++

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++) cout << v[i] << ", ";
    cout << "]"<<endl;
}
class Solution {
public:
    vector<int> threeEqualParts(vector<int>& A){
        vector<int> ret(2, -1);
        int num = 0, n = A.size();
        for (int i = 0; i < n; i++) num += (A[i] == 1);
        if (num % 3 != 0) return ret;
        if (num == 0) return {0, 2};
        int req = num / 3;
        int idx = n - 1;
        for (int temp = 0; idx >= 0 && temp < req; idx--) temp += A[idx] == 1;
        idx++;
        int firstEnd = getIdx(A, 0, idx);
        if (firstEnd < 0) return ret;
        int secondEnd = getIdx(A, firstEnd + 1, idx);
        if (secondEnd < 0) return ret;
        return { firstEnd, secondEnd + 1 };
    }
    int getIdx(vector<int>& a, int left, int right){
        while (left < right && a[left] == 0) left++;
        while (right < (int)a.size()) {
            if (a[left] != a[right]) return -1;
            left++; right++;
        }
        return left - 1;
    }
};
int main(){
    Solution ob;
    vector<int> v = {0,1,0,1,1};
    print_vector(ob.threeEqualParts(v));
}

Đầu vào

{0,1,0,1,1}

Đầu ra

[1, 4, ]