题目描述

给你一个整数数组 nums。好子序列 的定义是：子序列中任意 两个 连续元素的绝对差 恰好 为 1。

子序列 是指可以通过删除某个数组的部分元素（或不删除）得到的数组，并且不改变剩余元素的顺序。

返回 nums 中所有 可能存在的 好子序列的 元素之和。

因为答案可能非常大，返回结果需要对 10^9 + 7 取余。

注意，长度为 1 的子序列默认为好子序列。

样例

输入：nums = [1,2,1]

输出：14

解释：

好子序列包括：[1], [2], [1], [1,2], [2,1], [1,2,1]。

这些子序列的元素之和为 14。

输入：nums = [3,4,5]

输出：40

解释：

好子序列包括：[3], [4], [5], [3,4], [4,5], [3,4,5]。

这些子序列的元素之和为 40。

限制

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

算法

(动态规划) $O(n + U)$

1. 设状态 $f(j)$ 表示以元素 $j$ 结尾的合法子序列的个数，$s(j)$ 表示以元素 $j$ 结尾的合法子序列的元素之和。
2. 初始时，$f(nums(0)) = 1$，$s(nums(0)) = nums(0)$。
3. 然后从位置 $1$ 开始遍历数组，对于 $nums(i)$，可以从 $nums(i) - 1$，即 $f(nums(i)) = f(nums(i)) + f(nums(i) - 1)$，$s(nums(i)) = s(nums(i)) + s(nums(i) - 1) + f(nums(i - 1)) * nums(i)$。从 $nums(i) + 1$ 转移同理。也可以自己单独新开一个子序列，即 $f(nums(i)) = f(nums(i)) + 1$，$s(nums(i)) = s(nums(i)) + nums(i)$。
4. 最终答案为 $\sum s(j)$。

时间复杂度

• 初始化数组的时间复杂度为 $O(U)$。其中 $U$ 表示最大可能的数字。
• 动态规划遍历数组一次。
• 故总时间复杂度为 $O(n + U)$。

空间复杂度

• 需要 $O(U)$ 的额外空间存储动态规划的状态。

C++ 代码

#define LL long long

const int N = 100010;

class Solution {
private:
    int f[N], s[N];

public:
    int sumOfGoodSubsequences(vector<int>& nums) {
        const int n = nums.size();
        const int mod = 1000000007;

        memset(f, 0, sizeof(f));
        memset(s, 0, sizeof(s));

        f[nums[0]] = 1;
        s[nums[0]] = nums[0];

        for (int i = 1; i < n; i++) {
            f[nums[i]] = (f[nums[i]] + f[nums[i] + 1]) % mod;
            s[nums[i]] = (s[nums[i]] +
                    s[nums[i] + 1] + (LL)(f[nums[i] + 1]) * nums[i]
                ) % mod;

            if (nums[i] > 0) {
                f[nums[i]] = (f[nums[i]] + f[nums[i] - 1]) % mod;
                s[nums[i]] = (s[nums[i]] +
                        s[nums[i] - 1] + (LL)(f[nums[i] - 1]) * nums[i]
                    ) % mod;
            }

            f[nums[i]] = (f[nums[i]] + 1) % mod;
            s[nums[i]] = (s[nums[i]] + nums[i]) % mod;
        }

        int ans = 0;
        for (int i = 0; i < N; i++)
            ans = (ans + s[i]) % mod;

        return ans;
    }
};