Skip to content

3284. Sum of Consecutive Subarrays

  • Time: $O(n)$
  • Space: $O(1)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution {
 public:
  int getSum(const std::vector<int>& nums) {
    const long sum = accumulate(nums.begin(), nums.end(), 0L) % kMod;
    return (getSum(nums, 1) + getSum(nums, -1) - sum) % kMod;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  // Returns the sum of all subarrays with a difference of `diff`.
  long getSum(const vector<int>& nums, int diff) {
    long res = nums[0];
    long summ = nums[0];
    long count = 1;
    for (int i = 1; i < nums.size(); ++i) {
      if (nums[i] == nums[i - 1] + diff) {
        count += 1;
        summ += count * nums[i];
      } else {
        count = 1;
        summ = nums[i];
      }
      res += summ;
      res %= kMod;
    }
    return res;
  };
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
  public int getSum(int[] nums) {
    final long sum = Arrays.stream(nums).asLongStream().sum() % MOD;
    return (int) ((getSum(nums, 1) + getSum(nums, -1) - sum) % MOD);
  }

  private static final int MOD = 1_000_000_007;

  // Returns the sum of all subarrays with a difference of `diff`.
  private long getSum(int[] nums, int diff) {
    long res = nums[0];
    long summ = nums[0];
    long count = 1;
    for (int i = 1; i < nums.length; ++i) {
      if (nums[i] == nums[i - 1] + diff) {
        count += 1;
        summ += count * nums[i];
      } else {
        count = 1;
        summ = nums[i];
      }
      res += summ;
      res %= MOD;
    }
    return res;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
  def getSum(self, nums: list[int]) -> int:
    def getSum(diff: int) -> int:
      """Returns the sum of all subarrays with a difference of `diff`."""
      res = nums[0]
      summ = nums[0]
      count = 1
      for prev, num in itertools.pairwise(nums):
        if num == prev + diff:
          count += 1
          summ += count * num
        else:
          count = 1
          summ = num
        res += summ
      return res

    return (getSum(1) + getSum(-1) - sum(nums)) % 1_000_000_007