变长子数组求和-LeetCode

3427. 变长子数组求和

🎯 问题描述

描述：
给你一个整数数组 nums。
对于数组中的每个元素 nums[i]，我们需要计算一个特定子数组的和。
该子数组以 i 结尾，其起始位置由 nums[i] 的值决定。
具体来说，对于下标 i，子数组的起始下标为 start = max(0, i - nums[i])。
你需要返回所有这些子数组和的总和。

说明：

• 这是一个前缀和的应用题。

• 核心策略：

  • 题目要求计算多个不同长度、不同起点的子数组和。
  • 如果直接对每个位置都遍历求和，时间复杂度会达到 $O(N^2)$。
  • 利用前缀和数组，我们可以在 $O(1)$ 的时间内计算任意区间的和，从而将总复杂度降为 $O(N)$。

• 逻辑定义：

  • 对于每个 i，计算区间 [max(0, i - nums[i]), i] 的元素和。
  • 最后将所有计算出的和累加。

• 示例：

• 示例 1：

输入：nums = [2, 1, 3]
输出：7
解释：
- i=0: nums[0]=2. start = max(0, 0-2) = 0. 区间 [0,0], 和为 2.
- i=1: nums[1]=1. start = max(0, 1-1) = 0. 区间 [0,1], 和为 2+1=3.
- i=2: nums[2]=3. start = max(0, 2-3) = 0. 区间 [0,2], 和为 2+1+3=6.
(注：此处仅为逻辑演示，具体数值取决于题目对 start 的精确定义)

💻 解题思路

思路1：前缀和 + 变长窗口

核心逻辑

1. 构建前缀和数组 pre：
   • pre[i+1] 表示 nums[0] 到 nums[i] 的累加和。
   • 通过一次遍历即可完成构建。
   • 注意：标准前缀和通常令 pre[0] = 0，这样 pre[i+1] 才是前 i+1 个元素的和。你的代码中 pre[0]=nums[0] 是一种变体，配合后续的减法逻辑使用。
2. 计算变长子数组和：
   • 遍历每个下标 i。
   • 确定起点：start = max(0, i - nums[i])。这意味着子数组的长度受限于当前元素的值。
   • 利用前缀和求和：区间 [start, i] 的和等于 pre[i+1] - pre[start]。
     • 这里利用了前缀和的性质：$Sum(L, R) = Prefix(R+1) - Prefix(L)$。
3. 累加结果：
   • 将每个位置计算出的子数组和累加到最终结果中。

代码实现

class Solution:
    def subarraySum(self, nums: List[int]) -> int:
        n = len(nums)
        if n == 0:
            return 0
            
        # 1. 构建前缀和数组
        # pre[i+1] 将存储 nums[0]...nums[i] 的和
        pre = [0] * (n + 1)        
        pre[0] = 0 
        for i, num in enumerate(nums):
            pre[i+1] = pre[i] + num 
        total_sum = 0
        
        # 2. 计算每个位置对应的变长子数组和
        for i, num in enumerate(nums):
            # 确定子数组的起始位置
            # 逻辑：起始位置 = 当前位置 - 当前值 (且不能小于0)
            start = max(0, i - nums[i])
            
            # 利用前缀和计算区间 [start, i] 的和
            # Sum = pre[i+1] - pre[start]
            current_sum = pre[i+1] - pre[start]
            
            total_sum += current_sum
            
        return total_sum

📊 性能分析

• 时间复杂度：$O(N)$
  • 构建前缀和数组：$O(N)$。
  • 遍历计算每个位置的子数组和：$O(N)$。
  • 每次查询区间和的时间是 $O(1)$。
  • 总体线性时间，效率很高。
• 空间复杂度：$O(N)$
  • 需要一个长度为 $N+1$ 的 pre 数组来存储前缀和。

思考

1. 前缀和的初始化：

   • 标准做法：pre[0] = 0。这样 pre[1] 才是 nums[0]，pre[i+1] 才是 nums[0] + ... + nums[i]。
   • 区间 [L, R] 的和公式为 pre[R+1] - pre[L]。

2. start = max(0, i - nums[i]) 的含义：

   • 这个公式定义了一个变长窗口。
   • 窗口的长度大致为 nums[i] + 1（如果 i >= nums[i]）。
   • 这通常出现在涉及“跳跃能力”或“影响范围”的题目中。

3. 为什么使用前缀和？

   • 如果不用前缀和，对于每个 i，我们都需要一个循环从 start 加到 i。
   • 最坏情况下（例如 nums[i] 很大），每次求和需要 $O(N)$，总复杂度变成 $O(N^2)$。
   • 前缀和通过预处理，用空间换取了时间，将查询操作降为 $O(1)$。

4. 空间优化：

   • 虽然前缀和需要 $O(N)$ 空间，但如果我们只需要最终总和，其实不需要存储整个 pre 数组。
   • 不过，由于 start 是跳变的（不是单调递增的滑动窗口），我们无法简单地用一个变量维护当前窗口和，因此保留前缀和数组是合理的。