1
class Solution {
2
  public int[] threeEqualParts(int[] arr) {
3
    List<Integer> ones = new ArrayList<>();
4
    for (int i = 0; i < arr.length; i++) {
5
      if (arr[i] == 1) {
6
        ones.add(i);
7
      }
8
    }
9
    if (ones.size() == 0) { // edge case
10
      return new int[] {0, 2};
11
    }
12
    int[] ans = new int[2];
13
    int each = ones.size() / 3;
14
    for (int i = 0; i < 2 && ones.size() % 3 == 0; i++) { // for the first 2 partitions
15
      for (int j = 0; j < each - 1; j++) { // compare gaps
16
        if (ones.get(j + 1 + i * each) - ones.get(j + i * each)
17
            != ones.get(j + 2 * each + 1) - ones.get(j + 2 * each)) return new int[] {-1, -1};
18
      }
19
      ans[i] =
20
          ones.get((i + 1) * each - 1)
21
              + i
22
              + (arr.length - 1 - ones.get(ones.size() - 1)); // cut point
23
    }
24
    return ones.size() % 3 > 0 || ans[0] >= ones.get(each) || ans[1] > ones.get(2 * each)
25
        ? new int[] {-1, -1}
26
        : ans;
27
  }
28
}