二进制数组三等分

1. 描述

二进制数组分割, 分割成三等份, 返回下标

2. 代码

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> splitArray(vector<int>& A) {
	
	int n = A.size();
	int cnt1 = accumulate(A.begin(), A.end(), 0);
	// int cnt1 = 0;
	// for(int i = 0; i < A.size(); i++){
	// 	cnt1 += A[i];
	// }

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

	if (!cnt1) return {0, 2};
	cnt1 /= 3;
	int cnt = 0, left[3], right[3];

	for (int i = 0; i < n; ++i) {
		if (!A[i]) continue;
		cnt++;
		if ((cnt-1)%cnt1 == 0) left[(cnt-1)/cnt1] = i;
		if (cnt%cnt1 == 0) right[(cnt-1)/cnt1] = i;
	}

	for (int i = 0; i < 3; ++i) {
		right[i] += n - 1 - right[2];
	}

	for (int i = 0; i <= right[0]-left[0]; ++i) {
		if (A[left[1]+i] != A[left[0]+i] || A[left[2]+i] != A[left[0]+i])
			return {-1, -1};
	}

	return {right[0], right[1]+1};
}

int main()
{	
	vector<int> arr{1, 0, 1, 0, 1};
	// vector<int> arr{1, 1, 0, 1, 1};

	// int n;
	// vector<int> arr;
	// while(cin >> n){
	// 	arr.push_back(n);
	// 	if(cin.get() == '\n') break; // ADD
	// }

	vector<int> res = splitArray(arr);

	for(int i=0; i < res.size() - 1; i++){
		cout << res[i] << ",";
	}

	cout << res[res.size() - 1];

	return 0;
}