算法刷题打卡第79天：生成平衡数组的方案数

生成平衡数组的方案数

难度：中等

给你一个整数数组 nums 。你需要选择 恰好 一个下标（下标从 0 开始）并删除对应的元素。请注意剩下元素的下标可能会因为删除操作而发生改变。

比方说，如果 nums = [6,1,7,4,1] ，那么：

选择删除下标 1 ，剩下的数组为 nums = [6,7,4,1] 。
选择删除下标 2 ，剩下的数组为 nums = [6,1,4,1] 。
选择删除下标 4 ，剩下的数组为 nums = [6,1,7,4] 。
如果一个数组满足奇数下标元素的和与偶数下标元素的和相等，该数组就是一个 平衡数组 。

请你返回删除操作后，剩下的数组 nums 是 平衡数组 的 方案数 。

示例 1：

输入：nums = [2,1,6,4]
输出：1
解释：
删除下标 0 ：[1,6,4] -> 偶数元素下标为：1 + 4 = 5 。奇数元素下标为：6 。不平衡。
删除下标 1 ：[2,6,4] -> 偶数元素下标为：2 + 4 = 6 。奇数元素下标为：6 。平衡。
删除下标 2 ：[2,1,4] -> 偶数元素下标为：2 + 4 = 6 。奇数元素下标为：1 。不平衡。
删除下标 3 ：[2,1,6] -> 偶数元素下标为：2 + 6 = 8 。奇数元素下标为：1 。不平衡。
只有一种让剩余数组成为平衡数组的方案。

示例 2：

输入：nums = [1,1,1]
输出：3
解释：你可以删除任意元素，剩余数组都是平衡数组。

示例 3：

输入：nums = [1,2,3]
输出：0
解释：不管删除哪个元素，剩下数组都不是平衡数组。

动态规划

思路：

首先题目给出一个下标从 $0$ 开始长度为 $n$ 的整数数组 $\text{nums}$，并给出「平衡数组」的定义：数组中奇数下标元素和与偶数下标元素的和相等。现在我们可以选择一个位置 $i$，$0\le i<n$ 的元素删除（注意删除后 $i$ 之后的下标可能会因此发生改变）。我们需要求出所有删除该下标元素后的数组为「平衡数组」的下标数目。

我们设 $\text{preOdd}[i]$ 表示位置 $i$，$0\le i<n$ 前所有奇数位置元素的和，$\text{preEven}[i]$ 表示位置 $i$ 前所有偶数位置元素的和，$\text{sufOdd}[i]$ 表示位置 $i$ 后所有奇数位置元素的和，$\text{sufEven}[i]$ 表示位置 $i$ 后所有偶数位置元素的和，我们来看如何进行状态转移——

• 当 $i$ 为奇数时：
  • $\text{preOdd}[i+1]=\text{preOdd}[i]+\text{nums}[i]，i+1<n$
  • $\text{sufOdd}[i]=\text{sufOdd}[i-1]-\text{nums}[i]，i-1\ge 0$
  • $\text{preEven}[i+1]=\text{preEven}[i]，i+1<n$
  • $\text{sufEven}[i]=\text{sufEven}[i-1]，i-1\ge 0$
• 当 $i$ 为偶数时：
  • $\text{preEven}[i+1]=\text{preEven}[i]+\text{nums}[i]，i+1<n$
  • $\text{sufEven}[i]=\text{sufEven}[i-1]-\text{nums}[i]，i-1<n$
  • $\text{preOdd}[i+1]=\text{preOdd}[i]，i+1<n$
  • $\text{sufOdd}[i]=\text{sufOdd}[i-1]，i-1\ge 0$

其中边界条件为：当 $i=0$ 时，$\text{preOdd}[0]=0$，$\text{preEven}[0]=0$，$\text{sufOdd}[0]={\sum }_{2i+1}^n\text{nums}[2i+1]$，$\text{sufEven}[0]={\sum }_{2i}^n\text{nums}[2i]$。

不失一般性，现在我们将下标 $i$ 的元素进行删除，显而易见下标 $i$ 之前的元素下标并不会因此发生改变，而下标 $i$ 之后的原本在 $j$，$j>i$ 下标的数组元素会移动到下标 $j-1$，即下标 $i$ 之后的奇数下标元素会成为偶数下标元素，偶数下标元素会成为奇数下标元素。所以删除后数组中全部奇数下标元素和为 $\text{preOdd}[i]+\text{sufEven}[i]$，全部偶数下标元素和为 $\text{preEven}[i]+\text{sufOdd}[i]$，若两者相等则说明删除下标 $i$ 后的数组为「平衡数组」。那么我们尝试删除每一个下标 $i$，$0\le i<n$，来统计能生成「平衡数组」的下标即可。又因为 $\text{preOdd}[i]$，$\text{preEven}[i]$，$\text{sufOdd}[i]$，$\text{sufEven}[i]$ 求解都与前一个数有关，因此我们可以用「滚动数组」的技巧来进行空间优化。

复杂度分析：

• 时间复杂度： $O(n)$，其中 $n$ 为数组 $\text{nums}$ 的长度。
• 空间复杂度： $O(1)$，仅使用常量空间。

class Solution:
    def waysToMakeFair(self, nums: List[int]) -> int:
        res = 0
        q_j, q_o, h_j, h_o = 0, 0, sum(nums[::2]), sum(nums[1::2])
        for i in range(len(nums)):
            if i % 2 == 0:
                h_j -= nums[i]
                if h_o + q_j == h_j + q_o:
                    res += 1
                q_j += nums[i]
            else:
                h_o -= nums[i]
                if h_o + q_j == h_j + q_o:
                    res += 1
                q_o += nums[i]
        return res

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/ways-to-make-a-fair-array