214小R的权值计算

题解

问题理解

给定一个长度为 n 的数组 a，小 R 定义任意子数组的权值为 1×b₁ + 2×b₂ + ... + m×bₘ，其中 m 是子数组的长度，b₁, b₂, ..., bₘ 是子数组中的元素。要求计算所有子数组权值的和，结果需对 10^9 + 7 取模。

解题思路

b1 b2 b3 b4 b5 b1+2b2 b2+2b3 b3+2b4 b4+2b5 b1+2b2+3b3 b2+2b3+3b4 b3+2b4+3b5 b1+2b2+3b3+4b4 b2+2b3+3b4+4b5 b1+2b2+3b3+4b4+5b5 根据规律可以看出长度为 length 的子数组转移的时候，只需要加上 length*a[i] - (s[i-1] - s[i-length-1])即可

1. 前缀和计算： 使用前缀和 s 来快速计算任意子数组的元素和，s[i] 表示前 i 个元素的和。

2. 权值计算： 对于每个子数组长度 length，遍历所有可能的子数组，计算当前子数组的权值 cur。

   • 初始时，计算第一个长度为 length 的子数组的权值。
   • 随着窗口的滑动，通过更新 cur 来计算下一个子数组的权值：

     cur += length * a[i] - (s[i - 1] - s[i - length - 1])
     

     这里，length * a[i] 是新加入元素对权值的贡献，(s[i - 1] - s[i - length - 1]) 是窗口内之前元素权值的调整。

代码实现

mod = int(1e9 + 7)
from itertools import accumulate

def solution(n: int, a: list) -> int:
    res = 0
    a = [0] + a
    s = list(accumulate(a))
    # 枚举子数组的长度
    for length in range(1, n + 1):
        cur = 0
        for i in range(1, length + 1):
            cur += a[i] * i
        res = (res + cur) % mod
        for i in range(length + 1, n + 1):
            cur += length * a[i] - (s[i - 1] - s[i - length - 1])
            res = (res + cur) % mod
    return res % mod

if __name__ == '__main__':
    print(solution(3, [1, 2, 3]) == 33)
    print(solution(4, [4, 5, 6, 7]) == 203)
    print(solution(2, [10, 20]) == 80)