Program to find out the sum of evenly spaced elements in an array in Python

Suppose, there is an array nums of size n containing positive integers. We have another array queries that contain integer pairs (pi, qi). For every query in the array queries, the answer will be the sum of numbers in the array nums[j] where pi ≤ j < n and (j - pi) is divisible by qi. We have to return the answer of all such queries, and if it is a large value, we return the answer modulo 10^9 + 7.

So, if the input is like nums = [2, 3, 4, 5, 6, 7, 8, 9, 10], queries = [(2, 5), (7, 3), (6, 4)], then the output will be [13, 9, 8].

Understanding the Problem

For each query (p, q), we need to find elements at positions starting from index p, then p+q, p+2q, p+3q, and so on, until we reach the end of the array. This creates an evenly spaced sequence with step size q.

Algorithm Steps

To solve this, we will follow these steps −

• A := nums

• Q := queries

• n := length of nums

• M := 10^9 + 7

• m := integer value of (n ^ 0.5) + 2

• P := a new list containing the list A m times

• for i in range 1 to m, do

  • for j in range n-1 to -1, decrease by 1, do

    • if i+j < n, then

      • P[i, j] := (P[i, j]+P[i, i+j]) modulo M

• for each value b, k in Q, do

  • if k < m, then

    • return [value from index P[k, b]]

  • otherwise

    • return [sum(A[b to k]) modulo M]

Example

Let us see the following implementation to get better understanding −

def solve(A, Q):
    n, M = len(A), 10**9+7
    m = int(n**0.5)+2
    P = [A[:] for _ in range(m)]
    
    for i in range(1, m):
        for j in range(n-1, -1, -1):
            if i+j < n:
                P[i][j] = (P[i][j] + P[i][i+j]) % M
    
    return [P[k][b] if k < m else sum(A[b::k]) % M for b, k in Q]

# Test the function
nums = [2, 3, 4, 5, 6, 7, 8, 9, 10]
queries = [(2, 5), (7, 3), (6, 4)]
result = solve(nums, queries)
print(result)

The output of the above code is −

[13, 9, 8]

Step-by-Step Example

Let's trace through the first query (2, 5) −

nums = [2, 3, 4, 5, 6, 7, 8, 9, 10]
# Query (2, 5): start at index 2, step by 5
# Positions: 2, 7
# Values: nums[2] + nums[7] = 4 + 9 = 13

# Query (7, 3): start at index 7, step by 3  
# Positions: 7
# Values: nums[7] = 9

# Query (6, 4): start at index 6, step by 4
# Positions: 6
# Values: nums[6] = 8

print("Query (2, 5):", nums[2] + nums[7])
print("Query (7, 3):", nums[7])
print("Query (6, 4):", nums[6])

Query (2, 5): 13
Query (7, 3): 9
Query (6, 4): 8

How It Works

The algorithm uses square root decomposition for optimization. For small step sizes (less than ?n), it precomputes prefix sums. For large step sizes, it calculates the sum directly using slicing with step notation A[b::k].

Conclusion

This solution efficiently handles queries for evenly spaced elements using square root decomposition. It precomputes results for small steps and calculates directly for large steps, achieving optimal time complexity.