LeetCode-in-Java

927. Three Equal Parts

Hard

You are given an array arr which consists of only zeros and ones, divide the array into three 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:

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: arr = [1,0,1,0,1]

Output: [0,3]

Example 2:

Input: arr = [1,1,0,1,1]

Output: [-1,-1]

Example 3:

Input: arr = [1,1,0,0,1]

Output: [0,2]

Constraints:

Solution

public class Solution {
    public int[] threeEqualParts(int[] arr) {
        int ones = 0;
        for (int num : arr) {
            ones += num;
        }
        if (ones == 0) {
            return new int[] {0, 2};
        } else if (ones % 3 != 0) {
            return new int[] {-1, -1};
        }
        ones /= 3;
        int index1 = -1;
        int index2 = -1;
        int index3 = -1;
        int totalOnes = 0;
        for (int i = 0; i < arr.length; i++) {
            if (arr[i] == 0) {
                continue;
            }
            totalOnes += arr[i];
            if (totalOnes == 1) {
                index1 = i;
            } else if (totalOnes == ones + 1) {
                index2 = i;
            } else if (totalOnes == 2 * ones + 1) {
                index3 = i;
            }
        }
        while (index3 < arr.length) {
            if (arr[index1] == arr[index3] && arr[index2] == arr[index3]) {
                ++index1;
                ++index2;
                ++index3;
            } else {
                return new int[] {-1, -1};
            }
        }
        return new int[] {index1 - 1, index2};
    }
}