Sum of special evenly spaced elements in array

[wanxingliu94]

You are given a 0-indexed integer array nums consisting of n non-negative integers.

You are also given an array queries, where queries[i] = [xi, yi]. The answer to the ith query is the sum of all nums[j] where xi <= j < n and (j - xi) is divisible by yi.

Return an array answer where answer.length == queries.length and answer[i] is the answer to the ith query modulo 10^9 + 7.

Constraints:

• n == nums.length
• 1 <= n <= 5 * 10^4
• 0 <= nums[i] <= 10^9
• 1 <= queries.length <= 1.5 * 10^5
• 0 <= xi < n
• 1 <= yi <= 5 * 10^4

This is quite a challenging question for me, so all what I am going to write about is really from the solution that I read from here.

Solution explained

Let us first set up our skeleton code:

class Solution {
    public int[] solve(int[] nums, int[][] queries) {
        int n = nums.length;
        int m = queries.length;
        int mod = (int) 1e9 + 7;
        int[] res = new int[m];
        //this is a threshold we will use and explain later
        int thres = (int) Math.sqrt(n);
}

Given a query, [a, b] we are looking for the sum of elements in nums of the following form:

nums[a] + nums[a + b] + nums[a + 2*b] + nums[a + 3 * b] + ...

, hence the name of the question. The first crucial idea of the problem is to divide b​ into two cases:

1. b is greater than the square root of n(when b is large, and I will try to explain how to pick the threshold at the end of this post):

   when b is super large, then there will be only a few terms in the sum above, so we can just brutally calculate the sum by traversing over all of the terms involved. One such calculation will cost us O(nums.length / b) = O(square root of n), and we have at most n such elements. In total we need O(square root of n * n)

    for(int i = 0 ;i < queries.length; i++){
               int a = queries[i][0];
               int b = queries[i][1];
               if(b >  thres){
                   for(int j = 0; a + j * b < n; j++ ){
                       res[i] = (res[i] + nums[a + j*b]) % mod;
                   }
               }
    }

2. b is less than or equal to the square root of n(when b is small):

   In this case, we first notice the following observation:

   query[a + b, b]:  nums[a + b] + nums[a + 2 * b] + nums[a + 3 * b] + ...
   query[a, b]    :  nums[a] + query[a+b, b]      
   

   So the idea is whenever we encounter query[a, b], where b has never appeared in the previous queries which we have already scanned, we calculate and store all results of query[c, b] where c < n using the recurrence relation above, i.e. query[c, b] = query[c + b, b] + nums[c]. We will introduce a variable querySum to cache the results

   //the maximum value of b is thres since that is the definition of situation 2 whereas the maximum possible value of a is constrained by the problem description to be at most n - 1;
   int[][] querySum = new int[n][thres + 1];
   for (int i = 0; i < querySum.length; i++) {
       Arrays.fill(querySum[i], -1);
   }
   for (int i = 0; i < queries.length; i++) {
       int a = queries[i][0];
       int b = queries[i][1];
       if (b > thres) {
           for (int j = 0; a + j * b < n; j++) {
               res[i] = (res[i] + nums[a + j * b]) % mod;
           }
       } else {
           //If we have never scanned such b value before
           if (querySum[0][b] == -1) {
               for (int j = n - 1; j >= 0; j--) {
                   if (j + b >= n) {
                       querySum[j][b] = nums[j];
                   } else {
                       querySum[j][b] = (nums[j] + querySum[j + b][b]) % mod;
                   }
               }
           }
           res[i] = querySum[a][b];
   
       }
   }

   We will at most have have square root of n elements whose b is greater than the threshold, and for each memorization, it will cost us n operation. Thus in total for this situation we need to perform O(n * square root of n) again.

In summary, the time complexity of this algorithm is O(n * square root of n)​;

How to pick the threshold

If you are not familiar with single variable calculus, you might want to skip this. Now why do we choose the threshold to be the square root of n. If you go over the same thinking process using a random thres, then the time complexity will be O(n * n / thres + n * thres), so this becomes a simple optimization problem. Let us define

f(x) = n ^2 /  x + n * x

, then

f'(x) = - (n / x)^2 + n

. You can easily see that the critical point of this function is at x = the square root of n whatever the critical point means, and this concludes the content of this post.

References:

Github: https://github.com/wisdompeak/LeetCode/blob/master/Others/1714.Sum-Of-Special-Evenly-Spaced-Elements-In-Array/1714.Sum-Of-Special-Evenly-Spaced-Elements-In-Array_v2.cpp