1664. Ways to Make a Fair Array

1664. Ways to Make a Fair Array (Medium)

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

Related Topics:
Dynamic Programming, Greedy

Solution 1. Suffix Sum

Intuition:

If we remove A[i], then the parity of the indexes greater than i are all changed (even becomes odd and vice versa).

[ left part ] A[i] [ right part ]

We can get the sum of even/odd numbers in the right part using the precomputed suffix sum, and get the sum of even/odd numbers in the left part by calculating prefix sum.

Algorithm:

Let e[i]/o[i] be the (suffix) sum of all even/odd numbers with index greater than and equal to i.

We precompute e[i] and o[i] first.

Then we can scan from left to right and compute the prefix sum of all even/odd numbers with indexes smaller than i on the fly, and save them in even and odd.

If we remove A[i], then the sum of all the even numbers is even + o[i + 1], and the sum of all odd numbers is odd + e[i + 1].

// OJ: https://leetcode.com/problems/ways-to-make-a-fair-array/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    int waysToMakeFair(vector<int>& A) {
        int N = A.size(), even = 0, odd = 0, ans = 0;
        vector<int> e(N + 1), o(N + 1);
        for (int i = N - 1; i >= 0; --i) {
            if (i % 2 == 0) e[i] += A[i];
            else o[i] += A[i];
            e[i] += e[i + 1];
            o[i] += o[i + 1];
        }
        for (int i = 0; i < N; ++i) {
            ans += (even + o[i + 1]) == (odd + e[i + 1]);
            if (i % 2 == 0) even += A[i];
            else odd += A[i];
        }
        return ans;
    }
};

Solution 2.

We can compute the suffix sums on the fly as well to save the extra space

// OJ: https://leetcode.com/problems/ways-to-make-a-fair-array/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(1)
class Solution {
public:
    int waysToMakeFair(vector<int>& A) {
        int N = A.size(), evenPrefix = 0, oddPrefix = 0, evenSuffix = 0, oddSuffix = 0, ans = 0;
        for (int i = 0; i < N; ++i) {
            if (i % 2 == 0) evenSuffix += A[i];
            else oddSuffix += A[i];
        }
        for (int i = 0; i < N; ++i) {
            if (i % 2 == 0) evenSuffix -= A[i];
            else oddSuffix -= A[i];
            ans += (evenPrefix + oddSuffix) == (oddPrefix + evenSuffix);
            if (i % 2 == 0) evenPrefix += A[i];
            else oddPrefix += A[i];
        }
        return ans;
    }
};