927. 三等分

English Version

题目描述

给定一个由 0 和 1 组成的数组 arr ，将数组分成  3 个非空的部分 ，使得所有这些部分表示相同的二进制值。

如果可以做到，请返回任何 [i, j]，其中 i+1 < j，这样一来：

• arr[0], arr[1], ..., arr[i] 为第一部分；
• arr[i + 1], arr[i + 2], ..., arr[j - 1] 为第二部分；
• arr[j], arr[j + 1], ..., arr[arr.length - 1] 为第三部分。
• 这三个部分所表示的二进制值相等。

如果无法做到，就返回 [-1, -1]。

注意，在考虑每个部分所表示的二进制时，应当将其看作一个整体。例如，[1,1,0] 表示十进制中的 6，而不会是 3。此外，前导零也是被允许的，所以 [0,1,1] 和 [1,1] 表示相同的值。

 

示例 1：

输入：arr = [1,0,1,0,1]
输出：[0,3]

示例 2：

输入：arr = [1,1,0,1,1]
输出：[-1,-1]

示例 3:

输入：arr = [1,1,0,0,1]
输出：[0,2]

 

提示：

• 3 <= arr.length <= 3 * 104
• arr[i] 是 0 或 1

解法

方法一：将 1 的数量三等分

将 $1$ 的数量三等分，找到每一部分的第一个 $1$，分别记为 $i$, $j$, $k$。

然后从 $i$, $j$, $k$ 开始往后同时遍历每一部分，判断三部分对应的值是否相等，是则继续遍历，直至 $k$ 到达 $arr$ 末尾。

遍历结束时，若 $k=n$，说明满足三等分，返回此时的 $[i-1,j]$ 作为答案。

时间复杂度 $O(n)$，其中 $n$ 表示 $arr$ 的长度。

Python3

class Solution:
    def threeEqualParts(self, arr: List[int]) -> List[int]:
        def find(cnt):
            s = 0
            for i, v in enumerate(arr):
                s += v
                if s == cnt:
                    return i
            return -1

        n = len(arr)
        cnt, mod = divmod(sum(arr), 3)
        if mod:
            return [-1, -1]
        if cnt == 0:
            return [0, n - 1]
        i = find(1)
        j = find(cnt + 1)
        k = find(cnt * 2 + 1)
        while k < n and arr[i] == arr[j] == arr[k]:
            i, j, k = i + 1, j + 1, k + 1
        if k == n:
            return [i - 1, j]
        return [-1, -1]

Java

class Solution {
    public int[] threeEqualParts(int[] arr) {
        int n = arr.length;
        int cnt1 = 0;
        for (int v : arr) {
            cnt1 += v;
        }
        int cnt = cnt1 / 3;
        int mod = cnt1 % 3;
        if (mod != 0) {
            return new int[]{-1, -1};
        }
        if (cnt == 0) {
            return new int[]{0, n - 1};
        }
        int i = find(arr, 1);
        int j = find(arr, cnt + 1);
        int k = find(arr, cnt * 2 + 1);
        while (k < n && arr[i] == arr[j] && arr[j] == arr[k]) {
            ++i;
            ++j;
            ++k;
        }
        if (k == n) {
            return new int[]{i - 1, j};
        }
        return new int[]{-1, -1};
    }

    private int find(int[] arr, int cnt) {
        int s = 0;
        for (int i = 0; i < arr.length; ++i) {
            s += arr[i];
            if (s == cnt) {
                return i;
            }
        }
        return -1;
    }
}

C++

class Solution {
public:
    vector<int> threeEqualParts(vector<int>& arr) {
        int n = arr.size();
        int cnt1 = accumulate(arr.begin(), arr.end(), 0);
        int cnt = cnt1 / 3;
        int mod = cnt1 % 3;
        if (mod) return {-1, -1};
        if (cnt == 0) return {0, n - 1};
        int i = find(arr, 1);
        int j = find(arr, cnt + 1);
        int k = find(arr, cnt * 2 + 1);
        while (k < n && arr[i] == arr[j] && arr[j] == arr[k])
        {
            ++i;
            ++j;
            ++k;
        }
        if (k == n) return {i - 1, j};
        return {-1, -1};
    }

    int find(vector<int>& arr, int cnt) {
        int s = 0;
        for (int i = 0; i < arr.size(); ++i)
        {
            s += arr[i];
            if (s == cnt) return i;
        }
        return -1;
    }
};

Go

func threeEqualParts(arr []int) []int {
	n := len(arr)
	cnt1 := 0
	for _, v := range arr {
		cnt1 += v
	}
	cnt := cnt1 / 3
	mod := cnt1 % 3
	if mod != 0 {
		return []int{-1, -1}
	}
	if cnt == 0 {
		return []int{0, n - 1}
	}
	find := func(cnt int) int {
		s := 0
		for i, v := range arr {
			s += v
			if s == cnt {
				return i
			}
		}
		return -1
	}
	i, j, k := find(1), find(cnt+1), find(cnt*2+1)
	for k < n && arr[i] == arr[j] && arr[j] == arr[k] {
		i++
		j++
		k++
	}
	if k == n {
		return []int{i - 1, j}
	}
	return []int{-1, -1}
}

...