Problem

You are given an integer array nums. A good subsequence is defined as a subsequence of nums where the absolute difference between any two consecutive elements in the subsequence is exactly 1.

Return the sum of all possible good subsequences of nums.

Since the answer may be very large, return it modulo 10^9 + 7.

Note that a subsequence of size 1 is considered good by definition.

Examples

Example 1

1
2
3
4
5
6
7
8
9

Input: nums = [1,2,1]

Output: 14

Explanation:

  * Good subsequences are: `[1]`, `[2]`, `[1]`, `[1,2]`, `[2,1]`, `[1,2,1]`.
  * The sum of elements in these subsequences is 14.

Example 2

1
2
3
4
5
6
7
8
9

Input: nums = [3,4,5]

Output: 40

Explanation:

  * Good subsequences are: `[3]`, `[4]`, `[5]`, `[3,4]`, `[4,5]`, `[3,4,5]`.
  * The sum of elements in these subsequences is 40.

Constraints

  • 1 <= nums.length <= 10^5
  • 0 <= nums[i] <= 10^5

Solution

Approach

Let cnt[x] be the count of x in nums. For each value x, let dp[x] be the sum of all good subsequences ending with x. For each x in nums, dp[x] = cnt[x] * (dp[x-1] + dp[x+1] + 1), since you can append x to any good subsequence ending with x-1 or x+1, or start a new subsequence with x. The answer is the sum of dp[x] * x for all x.


Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
#include <vector>
#include <unordered_map>
using namespace std;
class Solution {
public:
    int sumOfGoodSubsequences(vector<int>& nums) {
        const int MOD = 1e9+7, N = 1e5+2;
        vector<long long> cnt(N), dp(N);
        for (int x : nums) cnt[x]++;
        long long ans = 0;
        for (int x = 0; x < N; ++x) {
            if (cnt[x] == 0) continue;
            dp[x] = (cnt[x] * ((x > 0 ? dp[x-1] : 0) + (x+1 < N ? dp[x+1] : 0) + 1)) % MOD;
            ans = (ans + dp[x] * x) % MOD;
        }
        return (int)ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution {
    public int sumOfGoodSubsequences(int[] nums) {
        int MOD = 1_000_000_007, N = 100_002;
        long[] cnt = new long[N], dp = new long[N];
        for (int x : nums) cnt[x]++;
        long ans = 0;
        for (int x = 0; x < N; ++x) {
            if (cnt[x] == 0) continue;
            dp[x] = cnt[x] * ((x > 0 ? dp[x-1] : 0) + (x+1 < N ? dp[x+1] : 0) + 1) % MOD;
            ans = (ans + dp[x] * x) % MOD;
        }
        return (int)ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution {
    fun sumOfGoodSubsequences(nums: IntArray): Int {
        val MOD = 1_000_000_007L
        val N = 100_002
        val cnt = LongArray(N)
        val dp = LongArray(N)
        for (x in nums) cnt[x]++
        var ans = 0L
        for (x in 0 until N) {
            if (cnt[x] == 0L) continue
            dp[x] = cnt[x] * ((if (x > 0) dp[x-1] else 0L) + (if (x+1 < N) dp[x+1] else 0L) + 1L) % MOD
            ans = (ans + dp[x] * x) % MOD
        }
        return ans.toInt()
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class Solution:
    def sumOfGoodSubsequences(self, nums: list[int]) -> int:
        MOD = 10**9+7
        N = 10**5+2
        cnt = [0]*(N)
        dp = [0]*(N)
        for x in nums:
            cnt[x] += 1
        ans = 0
        for x in range(N):
            if cnt[x] == 0:
                continue
            dp[x] = cnt[x] * ((dp[x-1] if x > 0 else 0) + (dp[x+1] if x+1 < N else 0) + 1) % MOD
            ans = (ans + dp[x] * x) % MOD
        return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
impl Solution {
    pub fn sum_of_good_subsequences(nums: Vec<i32>) -> i32 {
        const MOD: i64 = 1_000_000_007;
        const N: usize = 100_002;
        let mut cnt = vec![0i64; N];
        let mut dp = vec![0i64; N];
        for &x in &nums { cnt[x as usize] += 1; }
        let mut ans = 0i64;
        for x in 0..N {
            if cnt[x] == 0 { continue; }
            dp[x] = cnt[x] * ((if x > 0 { dp[x-1] } else { 0 }) + if x+1 < N { dp[x+1] } else { 0 } + 1) % MOD;
            ans = (ans + dp[x] * (x as i64)) % MOD;
        }
        ans as i32
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
function sumOfGoodSubsequences(nums: number[]): number {
    const MOD = 1_000_000_007, N = 100_002;
    const cnt = new Array(N).fill(0);
    const dp = new Array(N).fill(0);
    for (const x of nums) cnt[x]++;
    let ans = 0;
    for (let x = 0; x < N; ++x) {
        if (cnt[x] === 0) continue;
        dp[x] = cnt[x] * ((x > 0 ? dp[x-1] : 0) + (x+1 < N ? dp[x+1] : 0) + 1) % MOD;
        ans = (ans + dp[x] * x) % MOD;
    }
    return ans;
}

Complexity

  • ⏰ Time complexity: O(n + m) where n = len(nums), m = max(nums)
  • 🧺 Space complexity: O(m)