二进制数组三等分

• • 1. 描述
  • 2. 代码

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

#include 
#include 
#include 

using namespace std;

vector splitArray(vector& 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 arr{1, 0, 1, 0, 1};
	// vector arr{1, 1, 0, 1, 1};

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

	vector res = splitArray(arr);

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

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

	return 0;
}