python-leetcode-1664. 生成平衡数组的方案数

1664. 生成平衡数组的方案数 - 力扣（LeetCode）

可以使用前缀和的方法来高效计算删除每个元素后是否满足平衡数组的条件。核心思路如下：

1. 定义前缀和：

   • 计算 odd[i]：表示 nums[0] 到 nums[i] 之间 奇数下标 的元素和。
   • 计算 even[i]：表示 nums[0] 到 nums[i] 之间 偶数下标 的元素和。

2. 计算删除后的奇偶前缀和：

   • 若删除 nums[i]，则 i 及之后的元素索引都会发生变化。
   • 需要计算删除 nums[i] 后的 新偶数下标之和 和 新奇数下标之和，并进行比较。

3. 遍历判断：

   • 如果删除 nums[i] 后的奇偶和相等，则计数加一。

代码实现：

def waysToMakeFair(nums):
    n = len(nums)
    odd_prefix = [0] * (n + 1)  # 记录到 i 位置的奇数下标前缀和
    even_prefix = [0] * (n + 1)  # 记录到 i 位置的偶数下标前缀和

    # 计算奇数和偶数下标前缀和
    for i in range(n):
        odd_prefix[i + 1] = odd_prefix[i]
        even_prefix[i + 1] = even_prefix[i]
        if i % 2 == 0:
            even_prefix[i + 1] += nums[i]
        else:
            odd_prefix[i + 1] += nums[i]

    count = 0
    for i in range(n):
        # 删除 nums[i] 之前的偶数和奇数和
        even_sum = even_prefix[i]
        odd_sum = odd_prefix[i]
        
        # 删除 nums[i] 之后的调整
        new_even_sum = odd_sum + (even_prefix[n] - even_prefix[i + 1])
        new_odd_sum = even_sum + (odd_prefix[n] - odd_prefix[i + 1])
        
        if new_even_sum == new_odd_sum:
            count += 1

    return count

时间复杂度分析

• 预计算 odd_prefix 和 even_prefix 需要 O(n)。
• 遍历 nums 进行判断也需要 O(n)。
• 总时间复杂度为 O(n)，适用于大规模数据。

示例

nums = [2, 1, 6, 4]
print(waysToMakeFair(nums))  # 输出: 1

解释：

• 删除 nums[0] 后 [1, 6, 4] 不是平衡数组。
• 删除 nums[1] 后 [2, 6, 4] 是平衡数组，满足条件。

这样，我们可以高效计算删除某个元素后能否形成平衡数组，并返回满足条件的删除方案数。