Formatted question description: https://leetcode.ca/all/1714.html

1714. Sum Of Special Evenly-Spaced Elements In Array

Level

Hard

Description

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] = [x_i, y_i]. The answer to the i-th query is the sum of all nums[j] where x_i <= j < n and (j - x_i) is divisible by y_i.

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

Example 1:

Input: nums = [0,1,2,3,4,5,6,7], queries = [[0,3],[5,1],[4,2]]

Output: [9,18,10]

Explanation: The answers of the queries are as follows:

1) The j indices that satisfy this query are 0, 3, and 6. nums[0] + nums[3] + nums[6] = 9

2) The j indices that satisfy this query are 5, 6, and 7. nums[5] + nums[6] + nums[7] = 18

3) The j indices that satisfy this query are 4 and 6. nums[4] + nums[6] = 10

Example 2:

Input: nums = [100,200,101,201,102,202,103,203], queries = [[0,7]]

Output: [303]

Constraints:

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

Solution

Let n be the length of nums and let sqrt be the square root of n. Create a 2D array dp of n rows and sqrt - 1 columns, where dp[x][y] represents the query result of [x, y] when y is less than sqrt. For i from n - 1 to 0 and for j from 1 to sqrt - 1, if i + j < n, then dp[i][j] = nums[i] + dp[i + j][j]. Otherwise, dp[i][j] = nums[i].

For the i-th query [x, y], if y < sqrt, then answer[i] = dp[x][y]. Otherwise, calculate the sum accordingly and let answer[i] be the sum. Finally, return answer.

class Solution {
    public int[] solve(int[] nums, int[][] queries) {
        final int MODULO = 1000000007;
        int length = nums.length;
        int sqrt = (int) Math.sqrt(length);
        int queriesCount = queries.length;
        int[] answer = new int[queriesCount];
        long[][] dp = new long[length][sqrt];
        for (int i = length - 1; i >= 0; i--) {
            for (int j = 1; j < sqrt; j++) {
                if (i + j < length)
                    dp[i][j] = (nums[i] + dp[i + j][j]) % MODULO;
                else
                    dp[i][j] = nums[i];
            }
        }
        for (int i = 0; i < queriesCount; i++) {
            int x = queries[i][0], y = queries[i][1];
            if (y >= sqrt) {
                long sum = 0;
                while (x < length) {
                    sum = (sum + nums[x]) % MODULO;
                    x += y;
                }
                answer[i] = (int) sum;
            } else
                answer[i] = (int) dp[x][y];
        }
        return answer;
    }
}

All Problems

All Solutions