LeetCode 927. 三等分

题目描述

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

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

$A[0], A[1], ..., A[i]$ 组成第一部分；
$A[i+1], A[i+2], ..., A[j-1]$ 作为第二部分；
$A[j], A[j+1], ..., A[A.length - 1]$ 是第三部分。
这三个部分所表示的二进制值相等。
如果无法做到，就返回 $[-1, -1]$。

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

示例 1

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

示例 2

输出：[1,1,0,1,1]
输出：[-1,-1]

题解

题目要求将一个 0-1 数组分割成不为空的三部分，要求每部分看作二进制串后，数值相等（忽略前导零）

要满足数值相等，需要有以下条件：

• 每部分 1 的个数相同
• 每部分后导零个数相同
• 每部分除去前后导零外的部分相同

因此，需要操作的第一步是获取几个关键的 $1$ 的位置：每部分的第一个和最后一个 $1$
将第 $i$ 部分的第一个和最后一个 $1$ 分别记为 $one[i][0]$ 和 $one[i][1]$

对于有 $count1$ 个 $1$ 的数组，每部分应该有 $count1 / 3$ 个 $1$，且第 $i$ 部分的第一个 $1$ 为 $(i-1) \times count1 / 3 + 1$，最后一个 $1$ 为 $i \times count1 / 3$。这样可以在 $O(n)$ 的时间复杂度内，获得这些关键的位置。

前两部分的后导零可以通过划分成下一部分的前导零从而被“删去”，因此第三部分的后导零决定了后导零的个数。可以通过 $n - one[3][1] - 1$ 获取后导零的个数

那么，可根据每部分最后一个 $1$ 的位置加上后导零个数，获得分割点。
$\begin{cases} i = one[1][1] + zero \\ j = one[2][1] + zero + 1 \end{cases}$

接下来，分别判断 $i$、$j$ 是否合法（是否在下一个区域的位置）。至此，前两个条件已满足，只需要判断最后一个条件：非前后导零部分是否相同，只需要判断各个分区关键位置的长度即内容即可。

时间复杂度为 $O(n)$，空间复杂度为 $O(1)$。

除去上述部分，还需要考虑一些特殊情况：

• 数组为空（数组长度不足 $3$）: 直接返回 $[-1, -1]$
• 数组内容全为 $0$: 直接返回 $[0, 2]$（前一步已经确保长度至少为 $3$）

代码

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* threeEqualParts(int* A, int ASize, int* returnSize){
    int *res = malloc(sizeof(int)*2);
    res[0] = -1; res[1] = -1; *returnSize = 2;

    // 数组无法分成非空的 3 份
    if (ASize < 3) return res;

    int count = 0;
    for (int i=0; i<ASize; ++i)
        if (A[i] == 1)
            ++count;

    // 1 的个数无法整除
    if (count % 3 != 0) return res;
    // 全是 0
    if (count == 0) {res[0] = 0; res[1] = 2; return res;}

    // 获取关键的 1 的位置
    int p[6] = {0, 0, 0, 0, 0, 0};
    int n1 = count / 3; // 每组中 1 的长度
    int temp = 0;
    for (int i=0; i<ASize; ++i) {
        if (A[i] == 1) {
            ++temp;
            if (temp == 1) p[0] = i;            // 第一组第一个 1
            if (temp == n1) p[1] = i;           // 第一组的最后一个 1
            if (temp == n1 + 1) p[2] = i;       // 第二组的第一个 1
            if (temp == n1 * 2) p[3] = i;       // 第二组的最后一个 1
            if (temp == n1 * 2 + 1) p[4] = i;   // 第三组的第一个 1
            p[5] = i;                           // 第三组的最后一个 1
        }
    }

    // 后导零个数
    int zero = ASize - p[5] - 1;
    
    // 确定 i 和 j 的位置
    int i = p[1] + zero;
    int j = p[3] + zero + 1;
    if (i >= p[2] || j > p[4] || i+1 >= j)
        return res;
    
    // 判断该位置分割的区间是否合法
    if (p[1] - p[0] != p[3] - p[2] || p[1] - p[0] != p[5] - p[4])
        return res;
    for (int i=0; i< p[1]-p[0]; ++i)
        if (A[p[0] + i] != A[p[2] + i] || A[p[0] + i] != A[p[4] + i]) 
            return res;

    res[0] = i; res[1] = j;
    return res;
}