生成平衡数组的方案数

标签: 数组 前缀和

难度: Medium

给你一个整数数组 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
解释:不管删除哪个元素,剩下数组都不是平衡数组。

 

提示:

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 104

Submission

运行时间: 137 ms

内存: 21.0 MB

class Solution:
    def waysToMakeFair(self, nums) -> int:
        odd=sum(nums[::2])
        even=sum(nums[1::2])
        left_odd=0
        left_even=0
        ans=0
        for i in range(len(nums)):
            if not i%2:
                odd-=nums[i]
            else:
                even-=nums[i]
            if left_odd+even==left_even+odd:
                ans+=1
            if not i%2:
                left_odd+=nums[i]
            else:
                left_even+=nums[i]
        return ans

Explain

本题解采用前缀和思想,首先计算整个数组的奇数下标元素和(odd)与偶数下标元素和(even)。然后遍历数组,模拟删除每个位置的元素,通过维护两个变量left_odd和left_even来记录当前位置左侧的奇数和偶数和。对于每个位置i,首先从总和中减去当前元素,然后比较左侧偶数和加上右侧奇数和是否等于左侧奇数和加上右侧偶数和,如果相等,则当前删除操作产生的数组是平衡的。之后更新左侧奇偶和。最后返回所有平衡数组的方案数。

时间复杂度: O(n)

空间复杂度: O(1)

class Solution:
    def waysToMakeFair(self, nums) -> int:
        odd = sum(nums[::2])  # 奇数下标元素总和
        even = sum(nums[1::2])  # 偶数下标元素总和
        left_odd = 0  # 当前位置左侧奇数下标元素和
        left_even = 0  # 当前位置左侧偶数下标元素和
        ans = 0  # 计数平衡数组的方案数
        for i in range(len(nums)):
            if not i % 2:
                odd -= nums[i]  # 从奇数和中删除当前元素
            else:
                even -= nums[i]  # 从偶数和中删除当前元素
            if left_odd + even == left_even + odd:
                ans += 1  # 如果左侧偶数和加右侧奇数和等于左侧奇数和加右侧偶数和
            if not i % 2:
                left_odd += nums[i]  # 更新左侧奇数和
            else:
                left_even += nums[i]  # 更新左侧偶数和
        return ans  # 返回方案数

Explore

在删除操作的模拟中,我们并不真正地改变数组结构或其元素的下标。我们只是在逻辑上考虑如果某个元素被删除,会如何影响到总的奇数和偶数和。通过维护odd和even两个变量来分别记录所有奇数下标和偶数下标元素的和,然后根据当前元素的下标性质(奇或偶),直接从相应的和中减去该元素。这样做是因为我们只是想知道在不同删除情况下的偶和奇和的状态,而非改变数组的物理结构或下标。

在判断平衡数组条件之前,需要确保left_odd和left_even反映的是删除当前元素之前的状态,即当前元素左侧的奇偶和。如果先更新left_odd和left_even,那么比较的就不是删除当前元素之前的状态,而是包含了当前元素的新状态,这将导致错误的平衡判断。因此,我们先进行平衡性判断,然后再更新这些值以反映下一个元素处理时的左侧奇偶和。

在数组中,下标从0开始计数,因此偶数下标实际上是0, 2, 4, ..., 而奇数下标是1, 3, 5, ...。在题目中定义奇数和odd为奇数下标元素的和,偶数和even为偶数下标元素的和。因此,当我们遇到一个奇数下标的元素时,这个元素事实上是被计算在偶数和even中的。所以从even中减去这个奇数下标的元素是为了维护正确的偶数和状态。

平衡条件的核心思想是在任何删除元素的操作后,剩余数组的奇数下标元素和应等于偶数下标元素和。当删除i位置的元素后,i左侧的奇数和偶数和分别为left_odd和left_even。i右侧的奇数和偶数和可以通过总奇偶和减去左侧奇偶和和当前元素得出。特别地,我们需要注意当前元素下标的奇偶性,来决定它属于奇数和还是偶数和。最终,我们希望左侧奇数和加上右侧偶数和等于左侧偶数和加上右侧奇数和,这样确保整个数组在删除某个元素后保持平衡。