Prefix and

Posted on Edited on In Waline:

Today, let’s talk about a simple but very clever algorithm problem: calculate a subarray with several sums of k.

Then I will exhaust all the sub-arrays, calculate their sum, and see whose sum is equal to k.

The key is, ** how to quickly get the sum of a subarray **, for example, give you an array’nums’, let you implement an interface’sum (i, j) ‘, this interface to return’nums [i… j]’ and will be called multiple times, how do you implement this interface?

Since the interface needs to be called multiple times, it is obviously impossible to traverse’nums [i… j] ‘every time. Is there a fast way to calculate’nums [i… j]’ in O (1) time? This requires ** prefixes and ** tricks.

What is a prefix and

The idea of prefix sum is like this. For a given array’nums’, we create an additional prefix and array for preprocessing:

1
2
3
4
5
6
int n = nums.length;
//prefixes and arrays
int[] preSum = new int[n + 1];
preSum[0] = 0;
for (int i = 0; i < n; i++)
    preSum[i + 1] = preSum[i] + nums[i];

The meaning of this prefix and the array’preSum ‘is also well understood.’ preSum [i] ‘is the sum of’nums [0… i-1]’. So if we want to find the sum of’nums [i… j] ‘, we only need to operate’preSum [j + 1] -preSum [i]’ in one step, without having to traverse the array again.

Returning to this subarray problem, we want to find the sum of how many subarrays there are, and it is easy to write a solution with the help of prefixes and tricks:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
int subarraySum(int[] nums, int k) {
    int n = nums.length;
    //construct prefix and
    int[] sum = new int[n + 1];
    sum[0] = 0; 
    for (int i = 0; i < n; i++)
        sum[i + 1] = sum[i] + nums[i];
    
    int ans = 0;
    //Execute all subarrays
    for (int i = 1; i <= n; i++)
        for (int j = 0; j < i; j++)
            // sum of nums[j..i-1]
            if (sum[i] - sum[j]  k)
                ans++;

    return ans;
}

The time complexity $O (N ^ 2) $space complexity $O (N) $of this solution is not the optimal solution. However, after understanding the working principle of prefixes and arrays through this solution, you can use some clever methods to further reduce the time complexity.

II. Optimized solution

The previous solution has nested for loops:

1
2
3
4
for (int i = 1; i <= n; i++)
    for (int j = 0; j < i; j++)
        if (sum[i] - sum[j]  k)
            ans++;

What is the second layer of for loop doing? Translated, ** is calculating that there are several j’s that make the difference between sum [i] and sum [j] k. ** Whenever such a j 'is found, add the result by one.

We can shift the conditional judgment in the if statement and write it like this:

1
2
if (sum[j]  sum[i] - k)
    ans++;

The idea of optimization is: ** I directly record that there are several’sum [j] ‘and’sum [i] -k’ are equal, and update the result directly, which avoids the inner for loop **. We can use the hash table to record the number of occurrences of the prefix and the sum at the same time.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
int subarraySum(int[] nums, int k) {
    int n = nums.length;
    //map: prefix and - > the prefix and the number of occurrences
    HashMap<Integer, Integer> 
        preSum = new HashMap<>();
    // base case
    preSum.put(0, 1);

    int ans = 0, sum0_i = 0;
    for (int i = 0; i < n; i++) {
        sum0_i += nums[i];
        //this is the prefix and nums we are looking for [0.. j]
        int sum0_j = sum0_i - k;
        //If there is this prefix and in front, update the answer directly
        if (preSum.containsKey(sum0_j))
            ans += preSum.get(sum0_j);
        //add the prefix and nums [0.. i] and record the number of occurrences
        preSum.put(sum0_i, 
            preSum.getOrDefault(sum0_i, 0) + 1);
    }
    return ans;
}

Reprinted from:

https://github.com/labuladong/fucking-algorithm/blob/master/%E7%AE%97%E6%B3%95%E6%80%9D%E7%BB%B4%E7%B3%BB%E5%88%97/%E5%89%8D%E7%BC%80%E5%92%8C%E6%8A%80%E5%B7%A7.md