Question

1664. Ways to Make a Fair Array

中文文档

Description

You are given an integer array nums. You can choose exactly one index (0-indexed) and remove the element. Notice that the index of the elements may change after the removal.

For example, if nums = [6,1,7,4,1]:

• Choosing to remove index 1 results in nums = [6,7,4,1].
• Choosing to remove index 2 results in nums = [6,1,4,1].
• Choosing to remove index 4 results in nums = [6,1,7,4].

An array is fair if the sum of the odd-indexed values equals the sum of the even-indexed values.

Return the _number of indices that you could choose such that after the removal, _nums_ __is fair. _

 

Example 1:

Input: nums = [2,1,6,4]
Output: 1
Explanation:
Remove index 0: [1,6,4] -> Even sum: 1 + 4 = 5. Odd sum: 6. Not fair.
Remove index 1: [2,6,4] -> Even sum: 2 + 4 = 6. Odd sum: 6. Fair.
Remove index 2: [2,1,4] -> Even sum: 2 + 4 = 6. Odd sum: 1. Not fair.
Remove index 3: [2,1,6] -> Even sum: 2 + 6 = 8. Odd sum: 1. Not fair.
There is 1 index that you can remove to make nums fair.

Example 2:

Input: nums = [1,1,1]
Output: 3
Explanation: You can remove any index and the remaining array is fair.

Example 3:

Input: nums = [1,2,3]
Output: 0
Explanation: You cannot make a fair array after removing any index.

 

Constraints:

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

Solutions

Solution 1: Enumeration + Prefix Sum

First, we preprocess to get the sum $s_1$ of the elements at even indices and the sum $s_2$ of the elements at odd indices in the array nums.

Then, we enumerate each element $v$ in the array nums from front to back, using variables $t_1$ and $t_2$ to record the sum of the elements at even indices and the sum of the elements at odd indices that have been traversed, respectively.

We observe that for the current element $v$ being traversed, if it is removed, the sum of the elements at odd and even indices after this element will be swapped. At this point, we first determine whether the index $i$ of the current position is odd or even.

• If it is an even index, after deleting the element, the sum of the elements at odd indices in the array is $t_2 + s_1 - t_1 - v$, and the sum of the elements at even indices is $t_1 + s_2 - t_2$. If these two sums are equal, it is a balanced array, and the answer is incremented by one.

• If it is an odd index, after deleting the element, the sum of the elements at even indices in the array is $t_1 + s_2 - t_2 - v$, and the sum of the elements at odd indices is $t_2 + s_1 - t_1$. If these two sums are equal, it is a balanced array, and the answer is incremented by one.

Then we update $t_1$ and $t_2$ and continue to traverse the next element. After traversing the array, we can get the answer.

The time complexity is $O(n)$, where $n$ is the length of the array. The space complexity is $O(1)$.

class Solution:
  def waysToMakeFair(self, nums: List[int]) -> int:
    s1, s2 = sum(nums[::2]), sum(nums[1::2])
    ans = t1 = t2 = 0
    for i, v in enumerate(nums):
      ans += i % 2 == 0 and t2 + s1 - t1 - v == t1 + s2 - t2
      ans += i % 2 == 1 and t2 + s1 - t1 == t1 + s2 - t2 - v
      t1 += v if i % 2 == 0 else 0
      t2 += v if i % 2 == 1 else 0
    return ans

class Solution {
  public int waysToMakeFair(int[] nums) {
    int s1 = 0, s2 = 0;
    int n = nums.length;
    for (int i = 0; i < n; ++i) {
      s1 += i % 2 == 0 ? nums[i] : 0;
      s2 += i % 2 == 1 ? nums[i] : 0;
    }
    int t1 = 0, t2 = 0;
    int ans = 0;
    for (int i = 0; i < n; ++i) {
      int v = nums[i];
      ans += i % 2 == 0 && t2 + s1 - t1 - v == t1 + s2 - t2 ? 1 : 0;
      ans += i % 2 == 1 && t2 + s1 - t1 == t1 + s2 - t2 - v ? 1 : 0;
      t1 += i % 2 == 0 ? v : 0;
      t2 += i % 2 == 1 ? v : 0;
    }
    return ans;
  }
}

class Solution {
public:
  int waysToMakeFair(vector<int>& nums) {
    int s1 = 0, s2 = 0;
    int n = nums.size();
    for (int i = 0; i < n; ++i) {
      s1 += i % 2 == 0 ? nums[i] : 0;
      s2 += i % 2 == 1 ? nums[i] : 0;
    }
    int t1 = 0, t2 = 0;
    int ans = 0;
    for (int i = 0; i < n; ++i) {
      int v = nums[i];
      ans += i % 2 == 0 && t2 + s1 - t1 - v == t1 + s2 - t2;
      ans += i % 2 == 1 && t2 + s1 - t1 == t1 + s2 - t2 - v;
      t1 += i % 2 == 0 ? v : 0;
      t2 += i % 2 == 1 ? v : 0;
    }
    return ans;
  }
};

func waysToMakeFair(nums []int) (ans int) {
  var s1, s2, t1, t2 int
  for i, v := range nums {
    if i%2 == 0 {
      s1 += v
    } else {
      s2 += v
    }
  }
  for i, v := range nums {
    if i%2 == 0 && t2+s1-t1-v == t1+s2-t2 {
      ans++
    }
    if i%2 == 1 && t2+s1-t1 == t1+s2-t2-v {
      ans++
    }
    if i%2 == 0 {
      t1 += v
    } else {
      t2 += v
    }
  }
  return
}

/**
 * @param {number[]} nums
 * @return {number}
 */
var waysToMakeFair = function (nums) {
  let [s1, s2, t1, t2] = [0, 0, 0, 0];
  const n = nums.length;
  for (let i = 0; i < n; ++i) {
    if (i % 2 == 0) {
      s1 += nums[i];
    } else {
      s2 += nums[i];
    }
  }
  let ans = 0;
  for (let i = 0; i < n; ++i) {
    const v = nums[i];
    ans += i % 2 == 0 && t2 + s1 - t1 - v == t1 + s2 - t2;
    ans += i % 2 == 1 && t2 + s1 - t1 == t1 + s2 - t2 - v;
    t1 += i % 2 == 0 ? v : 0;
    t2 += i % 2 == 1 ? v : 0;
  }
  return ans;
};