2031. Count Subarrays With More Ones Than Zeros

一刷 06/2022

Version #1 Prefix Sum  Frequency Map

感觉最难的地方是理解为什么currGreater = prevGreater - currEqual

以及对于prefix sum永远要对 0)的情况做初始化

Time O(N)

Space O(N)

Runtime: 86 ms, faster than 23.75% of Java online submissions for Count Subarrays With More Ones Than Zeros.

Memory Usage: 109.7 MB, less than 45.00% of Java online submissions for Count Subarrays With More Ones Than Zeros.

 class Solution {

    private static int MOD = 1000000007;

    public int subarraysWithMoreZerosThanOnes(int[] nums) {

        int prefixSum = 0;

        int prevEqual = 0;

        int prevGreater = 0;

        // Ending with index i

        // dp[i][0] -- count of subarrays that have equal number of 0 and 1

        // dp[i][1] -- count of subarrays that have more 1s than 0s

        Map<Integer, Integer> prefixSumFreq = new HashMap<>();

        prefixSumFreq.put(0, 1);

        int result = 0;

        for (int i = 0; i < nums.length; i++) {

            if (nums[i] == 1) {

                prefixSum++;

            } else {

                prefixSum--;

            }

            int currEqual = prefixSumFreq.getOrDefault(prefixSum, 0);

            int currGreater = 0;

            if (nums[i] == 1) {

                currGreater = prevGreater + prevEqual + 1;

            } else {

                // Since currently those arrays have equal number of 0 and 1, before appending the

                // current 0, the count of 1 is just greater than count of 0 by 1

                currGreater = prevGreater - currEqual;

            }

            prefixSumFreq.put(prefixSum, prefixSumFreq.getOrDefault(prefixSum, 0) + 1);

            result = (currGreater + result) % MOD;

            prevEqual = currEqual;

            prevGreater = currGreater;

        }

        return result;

    }

}