Skip to content

3351. Sum of Good Subsequences 👍

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
 public:
  int sumOfGoodSubsequences(vector<int>& nums) {
    constexpr int kMod = 1'000'000'007;
    const int maxNum = ranges::max(nums);
    // endsIn[i] := the number of good subsequences ending in i
    vector<long> endsIn(maxNum + 3);
    // dp[i] := the sum of good subsequences ending in i
    vector<long> dp(maxNum + 3);

    for (const int num : nums) {
      const long seqsToAppend = 1 + endsIn[num] + endsIn[num + 2];
      dp[num + 1] =
          (seqsToAppend * num + (dp[num + 1] + dp[num] + dp[num + 2])) % kMod;
      endsIn[num + 1] = (endsIn[num + 1] + seqsToAppend) % kMod;
    }

    return accumulate(dp.begin(), dp.end(), 0,
                      [&](int acc, int d) { return (acc + d) % kMod; });
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
  public int sumOfGoodSubsequences(int[] nums) {
    final int MOD = 1000000007;
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    // endsIn[i] := the number of good subsequences ending in i
    long[] endsIn = new long[maxNum + 3];
    // dp[i] := the sum of good subsequences ending in i
    long[] dp = new long[maxNum + 3];

    for (final int num : nums) {
      final long seqsToAppend = 1 + endsIn[num] + endsIn[num + 2];
      dp[num + 1] = (seqsToAppend * num + (dp[num + 1] + dp[num] + dp[num + 2])) % MOD;
      endsIn[num + 1] = (endsIn[num + 1] + seqsToAppend) % MOD;
    }

    int ans = 0;
    for (final long d : dp)
      ans = (int) (ans + d) % MOD;
    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution:
  def sumOfGoodSubsequences(self, nums: list[int]) -> int:
    MOD = 10**9 + 7
    maxNum = max(nums)
    # endsIn[i] := the number of good subsequences ending in i
    endsIn = [0] * (maxNum + 2)
    # dp[i] := the sum of good subsequences ending in i
    dp = [0] * (maxNum + 2)

    for num in nums:
      seqsToAppend = 1 + endsIn[num - 1] + endsIn[num + 1]
      dp[num] = (seqsToAppend * num +
                 (dp[num] + dp[num - 1] + dp[num + 1])) % MOD
      endsIn[num] += seqsToAppend % MOD

    return sum(dp) % MOD