Skip to content

3299. Sum of Consecutive Subsequences 👍

  • Time: O(n)O(n)
  • Space: O(n)O(n)
 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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
class Solution {
 public:
  int getSum(vector<int>& nums) {
    const int increasingSequenceSum = getSequenceSum(nums, 1);
    const int decreasingSequenceSum = getSequenceSum(nums, -1);
    const int arraySum = getArraySum(nums);
    return (increasingSequenceSum + decreasingSequenceSum -
            static_cast<long>(arraySum) + kMod) %
           kMod;
  }

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

  // Returns the sum of all sequences in the array that are in consecutive
  // increasing order if `direction` is 1, or in consecutive decreasing order if
  // `direction` is -1.
  int getSequenceSum(const vector<int>& nums, int direction) {
    const int n = nums.size();
    long sequenceSum = 0;
    // {num: the number of subsequences ending in `num` so far}
    unordered_map<int, int> prefixCount;
    // {num: the number of subsequences starting from `num` so far}
    unordered_map<int, int> suffixCount;
    // prefixSubseqs[i] := the number of subsequences ending in nums[i]
    vector<int> prefixSubseqs(n);
    // suffixSubseqs[i] := the number of subsequences starting from nums[i]
    vector<int> suffixSubseqs(n);

    for (int i = 0; i < n; ++i) {
      const int prevNum = nums[i] - direction;
      const int subseqsCount =
          (prefixCount.contains(prevNum) ? prefixCount[prevNum] : 0) + 1;
      prefixSubseqs[i] = subseqsCount;
      prefixCount[nums[i]] += subseqsCount;
      prefixCount[nums[i]] %= kMod;
    }

    for (int i = n - 1; i >= 0; --i) {
      const int nextNum = nums[i] + direction;
      const int subseqsCount =
          (suffixCount.contains(nextNum) ? suffixCount[nextNum] : 0) + 1;
      suffixSubseqs[i] = subseqsCount;
      suffixCount[nums[i]] += subseqsCount;
      suffixCount[nums[i]] %= kMod;
    }

    for (int i = 0; i < n; ++i) {
      sequenceSum += static_cast<long>(nums[i]) * prefixSubseqs[i] % kMod *
                     suffixSubseqs[i];
      sequenceSum %= kMod;
    }

    return sequenceSum;
  }

  int getArraySum(const vector<int>& nums) {
    int arraySum = 0;
    for (const int num : nums)
      arraySum = (arraySum + num) % kMod;
    return arraySum;
  }
};

 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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class Solution {
  public int getSum(int[] nums) {
    final long increasingSequenceSum = getSequenceSum(nums, 1);
    final long decreasingSequenceSum = getSequenceSum(nums, -1);
    final long arraySum = getArraySum(nums);
    return (int) ((increasingSequenceSum + decreasingSequenceSum - arraySum + MOD) % MOD);
  }

  private static final int MOD = 1_000_000_007;

  // Returns the sum of all sequences in the array that are in consecutive
  // increasing order if `direction` is 1, or in consecutive decreasing order if
  // `direction` is -1.
  private long getSequenceSum(int[] nums, int direction) {
    int n = nums.length;
    long sequenceSum = 0;
    // {num: the number of subsequences ending in `num` so far}
    Map<Integer, Integer> prefixCount = new HashMap<>();
    // {num: the number of subsequences starting from `num` so far}
    Map<Integer, Integer> suffixCount = new HashMap<>();
    // prefixSubseqs[i] := the number of subsequences ending in nums[i]
    int[] prefixSubseqs = new int[n];
    // suffixSubseqs[i] := the number of subsequences starting from nums[i]
    int[] suffixSubseqs = new int[n];

    for (int i = 0; i < n; ++i) {
      final int prevNum = nums[i] - direction;
      final int subseqsCount = prefixCount.getOrDefault(prevNum, 0) + 1;
      prefixSubseqs[i] = subseqsCount;
      prefixCount.merge(nums[i], subseqsCount, Integer::sum);
      prefixCount.put(nums[i], prefixCount.get(nums[i]) % MOD);
    }

    for (int i = n - 1; i >= 0; --i) {
      int nextNum = nums[i] + direction;
      int subseqsCount = suffixCount.getOrDefault(nextNum, 0) + 1;
      suffixSubseqs[i] = subseqsCount;
      suffixCount.merge(nums[i], subseqsCount, Integer::sum);
      suffixCount.put(nums[i], suffixCount.get(nums[i]) % MOD);
    }

    for (int i = 0; i < n; ++i) {
      sequenceSum += (long) nums[i] * prefixSubseqs[i] % MOD * suffixSubseqs[i];
      sequenceSum %= MOD;
    }

    return sequenceSum;
  }

  private int getArraySum(int[] nums) {
    int arraySum = 0;
    for (int num : nums)
      arraySum = (arraySum + num) % MOD;
    return arraySum;
  }
}

 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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution:
  def getSum(self, nums: list[int]) -> int:
    MOD = 1_000_000_007
    n = len(nums)

    def getSequenceSum(nums: list[int], direction: int) -> int:
      """
      Returns the sum of all sequences in the array that are in consecutive
      increasing order if `direction` is 1, or in consecutive decreasing order
      if `direction` is -1."""
      sequenceSum = 0
      # {num: the number of subsequences ending in `num` so far}
      prefixCount = collections.Counter()
      # {num: the number of subsequences starting from `num` so far}
      suffixCount = collections.Counter()
      # prefixSubseqs[i] := the number of subsequences ending in nums[i]
      prefixSubseqs = [0] * n
      # suffixSubseqs[i] := the number of subsequences starting from nums[i]
      suffixSubseqs = [0] * n

      for i, num in enumerate(nums):
        prevNum = num - direction
        freq = prefixCount[prevNum] + 1
        prefixSubseqs[i] = freq
        prefixCount[num] += freq
        prefixCount[num] %= MOD

      for i, num in reversed(list(enumerate(nums))):
        nextNum = num + direction
        freq = suffixCount[nextNum] + 1
        suffixSubseqs[i] = freq
        suffixCount[num] += freq
        suffixCount[num] %= MOD

      for num, prefixSubseq, suffixSubseq in zip(
              nums, prefixSubseqs, suffixSubseqs):
        sequenceSum += num * prefixSubseq * suffixSubseq
        sequenceSum %= MOD

      return sequenceSum

    increasingSequenceSum = getSequenceSum(nums, 1)
    decreasingSequenceSum = getSequenceSum(nums, -1)
    return (increasingSequenceSum + decreasingSequenceSum - sum(nums) + MOD) % MOD
Was this page helpful?