Question

1714. 数组中特殊等间距元素的和

English Version

题目描述

给定一个索引从 0 开始的整数类型数组 nums ，包含 n 个非负整数。

另外给定一个（包含查询指令的）数组 queries ，其中 queries[i] = [xi, yi]。 第 i 个查询指令的答案是 nums[j] 中满足该条件的所有元素的和： xi <= j < n 且 (j - xi) 能被 yi 整除。

返回一个数组_ _answer，其中_  _answer.length == queries.length 且 answer[i] 是第 i 个查询指令的答案对 109 + 7 取模。

 

示例 1:

输入: nums = [0,1,2,3,4,5,6,7], queries = [[0,3],[5,1],[4,2]]
输出: [9,18,10]
解释: 每次查询的答案如下：
1) 符合查询条件的索引 j 有 0、 3 和 6。 nums[0] + nums[3] + nums[6] = 9
2) 符合查询条件的索引 j 有 5、 6 和 7。 nums[5] + nums[6] + nums[7] = 18
3) 符合查询条件的索引 j 有 4 和 6。 nums[4] + nums[6] = 10

示例 2:

输入: nums = [100,200,101,201,102,202,103,203], queries = [[0,7]]
输出: [303]

 

提示：

• n == nums.length
• 1 <= n <= 5 * 104
• 0 <= nums[i] <= 109
• 1 <= queries.length <= 1.5 * 105
• 0 <= xi < n
• 1 <= yi <= 5 * 104

解法

方法一：分块

这道题是一道比较典型的分块题目，对于步长较大的查询，我们可以直接暴力求解；对于步长较小的查询，我们可以预处理出每个位置的后缀和，然后直接查询。

本题中，我们将步长较大的查询的步长限制为 $\sqrt{n}$，这样就可以保证每个查询的时间复杂度为 $O(\sqrt{n})$。

我们定义一个二维数组 $suf$，其中 $suf[i][j]$ be 表示从位置 $j$ 开始，步长为 $i$ 的后缀和。那么对于每个查询 $[x, y]$，我们可以分为两种情况：

• 如果 $y \le \sqrt{n}$，那么我们可以直接查询 $suf[y][x]$；
• 如果 $y \gt \sqrt{n}$，那么我们可以直接暴力求解。

时间复杂度 $O((n + m) \times \sqrt{n})$，空间复杂度 $O(n \times \sqrt{n})$。其中 $n$ 是数组的长度，而 $m$ 是查询的个数。

class Solution:
  def solve(self, nums: List[int], queries: List[List[int]]) -> List[int]:
    mod = 10**9 + 7
    n = len(nums)
    m = int(sqrt(n))
    suf = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
      for j in range(n - 1, -1, -1):
        suf[i][j] = suf[i][min(n, j + i)] + nums[j]
    ans = []
    for x, y in queries:
      if y <= m:
        ans.append(suf[y][x] % mod)
      else:
        ans.append(sum(nums[x::y]) % mod)
    return ans

class Solution {
  public int[] solve(int[] nums, int[][] queries) {
    int n = nums.length;
    int m = (int) Math.sqrt(n);
    final int mod = (int) 1e9 + 7;
    int[][] suf = new int[m + 1][n + 1];
    for (int i = 1; i <= m; ++i) {
      for (int j = n - 1; j >= 0; --j) {
        suf[i][j] = (suf[i][Math.min(n, j + i)] + nums[j]) % mod;
      }
    }
    int k = queries.length;
    int[] ans = new int[k];
    for (int i = 0; i < k; ++i) {
      int x = queries[i][0];
      int y = queries[i][1];
      if (y <= m) {
        ans[i] = suf[y][x];
      } else {
        int s = 0;
        for (int j = x; j < n; j += y) {
          s = (s + nums[j]) % mod;
        }
        ans[i] = s;
      }
    }
    return ans;
  }
}

class Solution {
public:
  vector<int> solve(vector<int>& nums, vector<vector<int>>& queries) {
    int n = nums.size();
    int m = (int) sqrt(n);
    const int mod = 1e9 + 7;
    int suf[m + 1][n + 1];
    memset(suf, 0, sizeof(suf));
    for (int i = 1; i <= m; ++i) {
      for (int j = n - 1; ~j; --j) {
        suf[i][j] = (suf[i][min(n, j + i)] + nums[j]) % mod;
      }
    }
    vector<int> ans;
    for (auto& q : queries) {
      int x = q[0], y = q[1];
      if (y <= m) {
        ans.push_back(suf[y][x]);
      } else {
        int s = 0;
        for (int i = x; i < n; i += y) {
          s = (s + nums[i]) % mod;
        }
        ans.push_back(s);
      }
    }
    return ans;
  }
};

func solve(nums []int, queries [][]int) (ans []int) {
  n := len(nums)
  m := int(math.Sqrt(float64(n)))
  const mod int = 1e9 + 7
  suf := make([][]int, m+1)
  for i := range suf {
    suf[i] = make([]int, n+1)
    for j := n - 1; j >= 0; j-- {
      suf[i][j] = (suf[i][min(n, j+i)] + nums[j]) % mod
    }
  }
  for _, q := range queries {
    x, y := q[0], q[1]
    if y <= m {
      ans = append(ans, suf[y][x])
    } else {
      s := 0
      for i := x; i < n; i += y {
        s = (s + nums[i]) % mod
      }
      ans = append(ans, s)
    }
  }
  return
}

function solve(nums: number[], queries: number[][]): number[] {
  const n = nums.length;
  const m = Math.floor(Math.sqrt(n));
  const mod = 10 ** 9 + 7;
  const suf: number[][] = Array(m + 1)
    .fill(0)
    .map(() => Array(n + 1).fill(0));
  for (let i = 1; i <= m; ++i) {
    for (let j = n - 1; j >= 0; --j) {
      suf[i][j] = (suf[i][Math.min(n, j + i)] + nums[j]) % mod;
    }
  }
  const ans: number[] = [];
  for (const [x, y] of queries) {
    if (y <= m) {
      ans.push(suf[y][x]);
    } else {
      let s = 0;
      for (let i = x; i < n; i += y) {
        s = (s + nums[i]) % mod;
      }
      ans.push(s);
    }
  }
  return ans;
}