LeetCode 1714 - Sum Of Special Evenly-Spaced Elements In Array

The problem requires calculating sums of specific subsets of an integer array nums based on queries. Each query [xi, yi] specifies a starting index xi and a step yi. The sum for this query includes all elements nums[j] such that j starts at xi and increases in steps of yi (i.e.

leetcodehardarraydynamic-programming

LeetCode Problem 1714

Difficulty: 🔴 Hard
Topics: Array, Dynamic Programming

Solution

Problem Understanding

The problem requires calculating sums of specific subsets of an integer array nums based on queries. Each query [xi, yi] specifies a starting index xi and a step yi. The sum for this query includes all elements nums[j] such that j starts at xi and increases in steps of yi (i.e., j = xi, xi + yi, xi + 2*yi, ...) until the end of the array. The result for each query must be taken modulo 10^9 + 7.

The inputs are as follows: nums is a list of non-negative integers with length n up to 5 * 10^4. The queries array can contain up to 1.5 * 10^5 queries, and the step value yi in a query is at most 5 * 10^4. These constraints indicate that a naive solution iterating over all steps for every query could be too slow.

Key edge cases include queries where the step yi is larger than the remaining array length, queries that start at the last index, and arrays where all numbers are zero or large. The problem guarantees xi is valid and yi >= 1, so no need to check for invalid indices or zero steps.

Approaches

Brute Force

A brute-force approach would process each query independently by iterating over nums starting from xi and adding elements at intervals of yi. This works correctly, but with up to 1.5 * 10^5 queries and steps up to 5 * 10^4, the time complexity can reach O(Q * n / y) in the worst case, which is infeasible for the largest inputs.

Optimized Approach

The key observation is that queries with the same step yi can reuse precomputed sums for positions modulo yi. For a fixed yi, there are yi possible remainders when indexing into nums. By precomputing prefix sums for each remainder, any query can be answered in constant time. This is feasible because the maximum step yi is 5 * 10^4, limiting the number of sums we need to maintain.

The algorithm thus divides queries into two categories: queries with small yi (up to a threshold, e.g., √n) can be handled with precomputed sums, and queries with large yi are handled directly since there are fewer steps to sum. This hybrid approach balances time and space complexity efficiently.

Approach Time Complexity Space Complexity Notes
Brute Force O(Q * n / y) O(1) Simple iteration over each query, too slow for large inputs
Optimal O(n * √n + Q) O(n * √n) Precompute sums for small steps, handle large steps directly

Algorithm Walkthrough

  1. Define a threshold B as the square root of n. Steps yi below B are considered small, and steps above B are considered large. This threshold balances the number of precomputed sums and query computation.
  2. For all small steps yi <= B, create a dictionary where each key is a step and each value is a list of prefix sums for each possible remainder modulo yi.
  3. Precompute sums in reverse order for each remainder. For remainder r and step yi, prefix[r][i] contains the sum of nums[i], nums[i + yi], ... until the end of the array.
  4. For queries with small yi, retrieve the precomputed sum directly using the starting index xi modulo yi.
  5. For queries with large yi, iterate over the indices starting from xi with steps of yi and accumulate the sum directly, modulo 10^9 + 7.
  6. Collect all query results in the output array.

Why it works: Each element in nums is included in exactly the correct set for the sum by either the precomputed prefix sums for small steps or direct summation for large steps. By precomputing sums based on remainders, the algorithm ensures no repeated computation and maintains correctness for all queries.

Python Solution

from typing import List

class Solution:
    def solve(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        MOD = 10**9 + 7
        n = len(nums)
        B = int(n ** 0.5) + 1  # Threshold for small steps
        small = {}
        
        # Precompute sums for small steps
        for y in range(1, B):
            small[y] = [0] * y
            for i in range(n - 1, -1, -1):
                small[y][i % y] = (nums[i] + small[y][i % y]) % MOD
        
        answer = []
        for xi, yi in queries:
            if yi < B:
                answer.append(small[yi][xi % yi] % MOD)
            else:
                s = 0
                for j in range(xi, n, yi):
                    s = (s + nums[j]) % MOD
                answer.append(s)
        return answer

The implementation first defines a threshold B for small steps and precomputes sums for these steps. It then processes each query by either looking up a precomputed sum for small steps or directly summing elements for large steps. All sums are computed modulo 10^9 + 7.

Go Solution

package main

func solve(nums []int, queries [][]int) []int {
    const MOD = int(1e9 + 7)
    n := len(nums)
    B := 1
    for B*B <= n {
        B++
    }

    small := make(map[int][]int)
    for y := 1; y < B; y++ {
        small[y] = make([]int, y)
        for i := n - 1; i >= 0; i-- {
            r := i % y
            small[y][r] = (nums[i] + small[y][r]) % MOD
        }
    }

    answer := make([]int, len(queries))
    for idx, q := range queries {
        xi, yi := q[0], q[1]
        if yi < B {
            answer[idx] = small[yi][xi%yi]
        } else {
            s := 0
            for j := xi; j < n; j += yi {
                s = (s + nums[j]) % MOD
            }
            answer[idx] = s
        }
    }
    return answer
}

In Go, the approach is nearly identical. The map small stores precomputed sums for small steps. Care is taken with slice initialization and modulo operations to avoid overflow.

Worked Examples

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

Query Computation Result
[0,3] indices 0,3,6 → 0+3+6 9
[5,1] indices 5,6,7 → 5+6+7 18
[4,2] indices 4,6 → 4+6 10

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

Query Computation Result
[0,7] indices 0,7 → 100+203 303

The algorithm correctly computes sums by either looking up precomputed values or iterating directly.

Complexity Analysis

Measure Complexity Explanation
Time O(n * √n + Q) Precomputing small steps requires O(n*B), iterating large steps takes O(Q * n / B) ≤ O(Q)
Space O(n * √n) Storing precomputed sums for each remainder modulo y up to threshold B

The reasoning is that for small steps, the maximum number of sums precomputed is proportional to n * B, and each query is processed in O(1) for small steps or O(n / B) for large steps.

Test Cases

# Provided examples
assert Solution().solve([0,1,2,3,4,5,6,7], [[0,3],[5,1],[4,2]]) == [9,18,10]  # Example 1
assert Solution().solve([100,200,101,201,102,202,103,203], [[0,7]]) == [303]  # Example 2

# Edge cases
assert Solution().solve([0], [[0,1]]) == [0]  # Single element
assert Solution().solve([1], [[0,1]]) == [1]  # Single element non-zero
assert Solution().solve([1,2,3,4,5], [[4,10]]) == [5]  # Step larger than remaining array
assert Solution().solve([1,2,3,4,5], [[0,5]]) == [1]  # Step equals array length
assert Solution().solve([10]*10000, [[0,1],[0,10000