Problem statement

https://leetcode.com/problems/sum-of-special-evenly-spaced-elements-in-array/

Solution

In this problem we are asked to answer queries which are sums of elements with indexes being arithmetical progression. It is like cumulative sums, but with steps. Let us use the following approach:

  1. If step = y of progression is small enough, then we precalculate all sums. We choose small enough it it is <= sqrt(n), so we need O(n*sqrt(n)) time to calculate all sums.
  2. If step = y is big, then we directly calculate sum of elements in time no more than O(sqrt(n).

Complexity

Total time complexity is O((n+m)*sqrt(n)) to calculate all cumulative sums and to answer all queries. Space complexity is O(n*sqrt(n)).

Code

class Solution:
    def solve(self, nums, queries):
        MOD, n, m = 10**9 + 7, len(nums), len(queries)
        n_sqrt, ans = int(sqrt(n)), []
        dp = [[0]*n for _ in range(n_sqrt)]
        for j in range(1, n_sqrt):
            for i in range(n - 1, -1, -1):
                dp[j][i] = nums[i] + (dp[j][i + j] if i + j < n else 0)

        for i, (x, y) in enumerate(queries):
            ans += [sum(nums[x::y])] if y >= n_sqrt else [dp[y][x]]
        
        return ans