Skip to content

1714. Sum Of Special Evenly-Spaced Elements In Array

  • Time: $O(n\sqrt{n})$
  • Space: $O(n\sqrt{n})$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
 public:
  vector<int> solve(vector<int>& nums, vector<vector<int>>& queries) {
    constexpr int kMod = 1'000'000'007;
    const int n = nums.size();
    const int sqrtN = static_cast<int>(sqrt(n));
    vector<int> ans;
    // prefix[x][y] = sum(nums[x + ay]), where a >= 0 and x + ay < n
    vector<vector<int>> prefix(n, vector<int>(sqrtN));

    // Set prefix[i][j] to nums[i] to indicate the sequence starts with nums[i].
    for (int i = 0; i < n; ++i)
      for (int j = 0; j < sqrtN; ++j)
        prefix[i][j] = nums[i];

    for (int x = n - 1; x >= 0; --x)
      for (int y = 1; y < sqrtN; ++y)
        if (x + y < n) {
          prefix[x][y] += prefix[x + y][y];
          prefix[x][y] %= kMod;
        }

    for (const vector<int>& query : queries) {
      const int x = query[0];
      const int y = query[1];
      if (y < sqrtN) {
        ans.push_back(prefix[x][y]);
      } else {
        int sum = 0;
        for (int i = x; i < n; i += y)
          sum = (sum + nums[i]) % kMod;
        ans.push_back(sum);
      }
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
  public int[] solve(int[] nums, int[][] queries) {
    final int kMod = 1_000_000_007;
    final int n = nums.length;
    final int sqrtN = (int) Math.sqrt(n);
    int[] ans = new int[queries.length];
    // prefix[x][y] = sum(nums[x + ay]), where a >= 0 and x + ay < n
    int[][] prefix = new int[n][sqrtN];

    // Set prefix[i][j] to nums[i] to indicate the sequence starts with nums[i].
    for (int i = 0; i < n; ++i)
      for (int j = 0; j < sqrtN; ++j)
        prefix[i][j] = nums[i];

    for (int x = n - 1; x >= 0; --x)
      for (int y = 1; y < sqrtN; ++y)
        if (x + y < n) {
          prefix[x][y] += prefix[x + y][y];
          prefix[x][y] %= kMod;
        }

    for (int i = 0; i < queries.length; ++i) {
      final int x = queries[i][0];
      final int y = queries[i][1];
      if (y < sqrtN) {
        ans[i] = prefix[x][y];
      } else {
        int sum = 0;
        for (int j = x; j < n; j += y)
          sum = (sum + nums[j]) % kMod;
        ans[i] = sum;
      }
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
  def solve(self, nums: list[int], queries: list[list[int]]) -> list[int]:
    kMod = 10**9 + 7
    n = len(nums)
    sqrtN = int(n**0.5)
    # prefix[x][y] = sum(nums[x + ay]), where a >= 0 and x + ay < n
    # Set prefix[i][j] to nums[i] to indicate the sequence starts with nums[i].
    prefix = [[num] * sqrtN for num in nums]

    for x in range(n - 1, -1, -1):
      for y in range(1, sqrtN):
        if x + y < n:
          prefix[x][y] += prefix[x + y][y]
          prefix[x][y] %= kMod

    return [prefix[x][y] if y < sqrtN
            else sum(nums[x::y]) % kMod
            for x, y in queries]