寻找数组中满足和积相等条件的平衡索引位置（O(n)复杂度）

我来帮你搞定这个平衡索引的问题！你提到现有代码在数组不含0时能正常运行，但遇到0就失效，那我们来一步步解决这个问题，同时保证O(n)的时间复杂度。

现有代码的问题分析 #

你当前的代码在数组不含0时能正常工作，但遇到0就会出问题——因为右侧乘积一旦包含0，结果就会变成0，这时候需要检查左侧和是否为0，但原有逻辑可能没处理这种边界情况，或者在计算乘积时直接跳过了0的场景，导致漏判。

优化方案：前缀和+后缀乘积数组 #

要实现O(n)时间复杂度，我们可以通过预处理前缀和和后缀乘积来避免重复计算，具体步骤如下：

1. 预处理后缀乘积数组 #

从右往左遍历数组，计算每个位置开始到数组末尾的乘积。我们定义suffixProduct[i]为从索引i到N-1的元素乘积，同时约定空乘积（即右侧没有元素时）的值为1（比如当P是最后一个索引时，右侧没有元素，积为1）。

2. 预处理前缀和（边遍历边维护） #

从左往右遍历数组时，实时维护当前的前缀和（即0到当前索引前一位的元素和），这样不用额外存储整个前缀和数组，节省空间。

3. 遍历检查每个索引 #

对每个索引P，分别获取左侧和（P=0时为0，否则取当前维护的前缀和）和右侧积（P=N-1时为1，否则取suffixProduct[P+1]），比较两者是否相等，相等则记录该索引。

4. 处理边界情况 #

• 数组长度小于3时：比如长度为1，左右都无元素，和为0、积为1，不相等；长度为2时，只有当arr[1]=0（P=0时左侧和0等于右侧积0）或arr[0]=1（P=1时左侧和1等于右侧积1）才符合条件。
• 乘积溢出问题：如果数组元素过大，整数乘积会溢出，这时候可以用BigInteger来处理，避免数值错误。

完整Java代码示例 #

package com.array.balance.indexes;

import java.util.ArrayList;
import java.util.List;

public class FindBalancedIndexesInArray {
    public static void main(final String[] args) {
        int[] arr1 = {2, 3, 0, 1, 0};
        List<Integer> result1 = findBalancedIndexes(arr1);
        System.out.println("数组arr1的平衡索引：" + result1); // 输出：[0, 4]

        int[] arr2 = {0, -1, 0};
        List<Integer> result2 = findBalancedIndexes(arr2);
        System.out.println("数组arr2的平衡索引：" + result2); // 输出：[0, 2]
    }

    public static List<Integer> findBalancedIndexes(int[] arr) {
        List<Integer> balancedIndexes = new ArrayList<>();
        int n = arr.length;
        if (n == 0) {
            return balancedIndexes;
        }

        // 计算后缀乘积数组，用long减少溢出概率
        long[] suffixProduct = new long[n + 1];
        suffixProduct[n] = 1; // 空乘积定义为1
        for (int i = n - 1; i >= 0; i--) {
            suffixProduct[i] = arr[i] * suffixProduct[i + 1];
        }

        // 遍历数组，维护前缀和并检查平衡索引
        long prefixSum = 0;
        for (int p = 0; p < n; p++) {
            long leftSum = p == 0 ? 0 : prefixSum;
            long rightProduct = p == n - 1 ? 1 : suffixProduct[p + 1];

            if (leftSum == rightProduct) {
                balancedIndexes.add(p);
            }

            // 更新前缀和，加入当前元素
            prefixSum += arr[p];
        }

        return balancedIndexes;
    }
}

代码说明 #

• 用long类型存储和与乘积，减少溢出概率；如果数组元素极大，建议替换为BigInteger。
• 遍历一次计算后缀乘积，再遍历一次检查每个索引，总时间复杂度O(n)，空间复杂度O(n)（用于存储后缀乘积数组）。
• 完美处理含0的情况：比如当右侧包含0时，乘积为0，只要左侧和为0就会被判定为平衡索引。

空间优化（可选） #

如果想把空间复杂度降到O(1)，可以先计算总乘积，然后从左到右遍历维护左侧和，同时动态调整右侧乘积，但这种方法无法处理数组中有0的情况（因为除法会出错），所以还是推荐上面的后缀数组方法，在空间和正确性之间取平衡。

内容的提问来源于stack exchange，提问作者Do It