Leetcode 1524. 和为奇数的子数组数目 前缀和并优化时间到O(n)

O(n^2)的前缀和超时

const int mod = (int)1e9+7;
class Solution {
public:
    int numOfSubarrays(vector<int>& arr) {
        int n = arr.size(), res = 0;
        vector<int> preSum(n+1);
        for(int i=1;i<=n;i++){
            preSum[i] = preSum[i-1] + arr[i-1];
        }
        for(int i=0;i<n;i++){
            for(int j=i+1;j<=n;j++){
                int k = preSum[j] - preSum[i];
                if(k%2) res = (res+1)%mod;
            }
        }
        return res;
    }
};

 

这里时间可以优化到O(N)，优化的思路是在一次遍历的时候，维护前面前缀和为奇数的个数和前缀和为偶数的个数即可。

const int mod = (int)1e9+7;
class Solution {
public:
    int numOfSubarrays(vector<int>& arr) {
        int n = arr.size(), res = 0;
        vector<int> preSum(n+1);
        // 当前前缀和为奇数，减去当前前缀和为偶数的前缀和，就是奇数数组
        // 当前前缀和为偶数，减去当前前缀和为奇数的前缀和，就是偶数数组
        // 维护两个变量，表示之前前缀和为奇数以及偶数的个数
        int oddNum = 0, evenNum = 1;
        for(int i=1;i<=n;i++){
            preSum[i] = preSum[i-1] + arr[i-1];
            res = (res%mod + (preSum[i]%2?evenNum:oddNum)%mod)%mod;
            preSum[i]%2?oddNum++:evenNum++;
        }
        return res;
    }
};

空间也优化到O(n), 比较好的写法

const int mod = (int)1e9+7;
class Solution {
public:
    int numOfSubarrays(vector<int>& arr) {
        int n = arr.size(), res = 0, preSum = 0;
        int oddNum = 0, evenNum = 1;
        for(int i=1;i<=n;i++){
            preSum = (preSum+arr[i-1])%2;
            if(preSum){
                res = (res%mod+evenNum%mod)%mod;
                oddNum++;
            }else{
                res = (res%mod+oddNum%mod)%mod;
                evenNum++;
            }
        }
        return res;
    }
};