927. Three Equal Parts

Problem

Given an array A of 0s and 1s, divide the array into 3 non-empty parts such that all of these parts represent the same binary value.

If it is possible, return any [i, j] with i+1 < j, such that:

• A[0], A[1], ..., A[i] is the first part;
• A[i+1], A[i+2], ..., A[j-1] is the second part, and
• A[j], A[j+1], ..., A[A.length - 1] is the third part.
• All three parts have equal binary value.

If it is not possible, return [-1, -1].

Note that the entire part is used when considering what binary value it represents. For example, [1,1,0] represents 6 in decimal, not 3. Also, leading zeros are allowed, so [0,1,1] and [1,1] represent the same value.

Example 1:

Input: [1,0,1,0,1]
Output: [0,3]

Example 2:

Input: [1,1,0,1,1]
Output: [-1,-1]

Note:

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

Problem Summary

Given an array A consisting of 0s and 1s, divide the array into 3 non-empty parts such that all of these parts represent the same binary value. If it is possible, return any [i, j], where i+1 < j, such that:

• A[0], A[1], ..., A[i] form the first part;
• A[i+1], A[i+2], ..., A[j-1] are the second part;
• A[j], A[j+1], ..., A[A.length - 1] are the third part.
• The binary values represented by these three parts are equal.

If it is not possible, return [-1, -1].

Note that when considering the binary value represented by each part, it should be treated as a whole. For example, [1,1,0] represents 6 in decimal, not 3. In addition, leading zeros are allowed, so [0,1,1] and [1,1] represent the same value.

Constraints:

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

Solution Approach

• Given an array containing only 0s and 1s, find 2 split points so that the binary values of the 3 subarrays are exactly the same.
• The solution idea for this problem is not difficult; just simulate according to the problem statement. First count the number of 1s, total, then divide by 3 to get the number of 1s that should appear in each part. Some special cases need additional checks, such as when there are no 1s, in which case we can only split at the head and tail. If the number of 1s is not a multiple of 3, it also cannot be split to satisfy the problem requirements. Then find the index of the first 1, and use total/3 to find mid, the first split point. Then continue moving forward to find the second split point. After finding these 3 points, move these 3 points synchronously, and during the movement check whether the values corresponding to these 3 indices are equal. If they are all equal, and the last point can move to the end, then a solution satisfying the problem requirements has been found.

Code

package leetcode

func threeEqualParts(A []int) []int {
	n, ones, i, count := len(A), 0, 0, 0
	for _, a := range A {
		ones += a
	}
	if ones == 0 {
		return []int{0, n - 1}
	}
	if ones%3 != 0 {
		return []int{-1, -1}
	}
	k := ones / 3
	for i < n {
		if A[i] == 1 {
			break
		}
		i++
	}
	start, j := i, i
	for j < n {
		count += A[j]
		if count == k+1 {
			break
		}
		j++
	}
	mid := j
	j, count = 0, 0
	for j < n {
		count += A[j]
		if count == 2*k+1 {
			break
		}
		j++
	}
	end := j
	for end < n && A[start] == A[mid] && A[mid] == A[end] {
		start++
		mid++
		end++
	}
	if end == n {
		return []int{start - 1, mid}
	}
	return []int{-1, -1}
}