跳转至

3299. 连续子序列的和 🔒

题目描述

如果一个长度为 n 的数组 arr 符合下面其中一个条件，可以称它 连续：

  • 对于所有的 1 <= i < n，arr[i] - arr[i - 1] == 1。
  • 对于所有的 1 <= i < n，arr[i] - arr[i - 1] == -1。

数组的 值 是其元素的和。

例如，[3, 4, 5] 是一个值为 12 的连续数组，并且 [9, 8] 是另一个值为 17 的连续数组。而 [3, 4, 3] 和 [8, 6] 都不连续。

给定一个整数数组 nums，返回所有 连续 非空 子序列 的 值 之和。

由于答案可能很大，返回它对 109 + 7 取模 的值。

注意 长度为 1 的数组也被认为是连续的。

 

示例 1：

输入：nums = [1,2]

输出：6

解释：

连续子序列为 [1]，[2]，[1, 2]。

示例 2：

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

输出：31

解释：

连续子序列为：[1]，[4]，[2]，[3]，[1, 2]，[2, 3]，[4, 3]，[1, 2, 3]。

 

提示：

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 105

解法

方法一：枚举贡献

我们不妨统计每个元素 $\textit{nums}[i]$ 在多少个长度大于 $1$ 的连续子序列中出现，那么其个数乘以 $\textit{nums}[i]$ 就是 $\textit{nums}[i]$ 在所有长度大于 $1$ 的连续子序列中的贡献。我们将其累加，再加上所有元素的和，即为答案。

我们可以先统计连续递增子序列对答案的贡献，再统计连续递减子序列对答案的贡献，最后再加上所有元素的和即可。

在实现上，我们定义一个函数 $\textit{calc}(\textit{nums})$，其中 $\textit{nums}$ 是一个数组，返回 $\textit{nums}$ 所有长度大于 $1$ 的连续子序列的和。

在函数中，我们可以使用两个数组 $\textit{left}$ 和 $\textit{right}$ 分别记录每个元素 $\textit{nums}[i]$ 的左侧以 $\textit{nums}[i] - 1$ 结尾的连续递增子序列的个数，以及右侧以 $\textit{nums}[i] + 1$ 开头的连续递增子序列的个数。这样，我们就可以在 $O(n)$ 的时间复杂度内计算出 $\textit{nums}$ 在所有长度大于 $1$ 的连续子序列中的贡献。

在主函数中，我们首先调用 $\textit{calc}(\textit{nums})$ 计算出连续递增子序列对答案的贡献，然后将 $\textit{nums}$ 反转后再次调用 $\textit{calc}(\textit{nums})$ 计算出连续递减子序列对答案的贡献，最后再加上所有元素的和即为答案。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是数组 $\textit{nums}$ 的长度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution:
    def getSum(self, nums: List[int]) -> int:
        def calc(nums: List[int]) -> int:
            n = len(nums)
            left = [0] * n
            right = [0] * n
            cnt = Counter()
            for i in range(1, n):
                cnt[nums[i - 1]] += 1 + cnt[nums[i - 1] - 1]
                left[i] = cnt[nums[i] - 1]
            cnt = Counter()
            for i in range(n - 2, -1, -1):
                cnt[nums[i + 1]] += 1 + cnt[nums[i + 1] + 1]
                right[i] = cnt[nums[i] + 1]
            return sum((l + r + l * r) * x for l, r, x in zip(left, right, nums)) % mod

        mod = 10**9 + 7
        x = calc(nums)
        nums.reverse()
        y = calc(nums)
        return (x + y + sum(nums)) % mod

 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
class Solution {
    private final int mod = (int) 1e9 + 7;

    public int getSum(int[] nums) {
        long x = calc(nums);
        for (int i = 0, j = nums.length - 1; i < j; ++i, --j) {
            int t = nums[i];
            nums[i] = nums[j];
            nums[j] = t;
        }
        long y = calc(nums);
        long s = Arrays.stream(nums).asLongStream().sum();
        return (int) ((x + y + s) % mod);
    }

    private long calc(int[] nums) {
        int n = nums.length;
        long[] left = new long[n];
        long[] right = new long[n];
        Map<Integer, Long> cnt = new HashMap<>();
        for (int i = 1; i < n; ++i) {
            cnt.merge(nums[i - 1], 1 + cnt.getOrDefault(nums[i - 1] - 1, 0L), Long::sum);
            left[i] = cnt.getOrDefault(nums[i] - 1, 0L);
        }
        cnt.clear();
        for (int i = n - 2; i >= 0; --i) {
            cnt.merge(nums[i + 1], 1 + cnt.getOrDefault(nums[i + 1] + 1, 0L), Long::sum);
            right[i] = cnt.getOrDefault(nums[i] + 1, 0L);
        }
        long ans = 0;
        for (int i = 0; i < n; ++i) {
            ans = (ans + (left[i] + right[i] + left[i] * right[i] % mod) * nums[i] % mod) % mod;
        }
        return ans;
    }
}

 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
class Solution {
public:
    int getSum(vector<int>& nums) {
        using ll = long long;
        const int mod = 1e9 + 7;
        auto calc = [&](const vector<int>& nums) -> ll {
            int n = nums.size();
            vector<ll> left(n), right(n);
            unordered_map<int, ll> cnt;

            for (int i = 1; i < n; ++i) {
                cnt[nums[i - 1]] += 1 + cnt[nums[i - 1] - 1];
                left[i] = cnt[nums[i] - 1];
            }

            cnt.clear();

            for (int i = n - 2; i >= 0; --i) {
                cnt[nums[i + 1]] += 1 + cnt[nums[i + 1] + 1];
                right[i] = cnt[nums[i] + 1];
            }

            ll ans = 0;
            for (int i = 0; i < n; ++i) {
                ans = (ans + (left[i] + right[i] + left[i] * right[i] % mod) * nums[i] % mod) % mod;
            }
            return ans;
        };

        ll x = calc(nums);
        reverse(nums.begin(), nums.end());
        ll y = calc(nums);
        ll s = accumulate(nums.begin(), nums.end(), 0LL);
        return static_cast<int>((x + y + s) % mod);
    }
};

 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
func getSum(nums []int) int {
    const mod = 1e9 + 7

    calc := func(nums []int) int64 {
        n := len(nums)
        left := make([]int64, n)
        right := make([]int64, n)
        cnt := make(map[int]int64)

        for i := 1; i < n; i++ {
            cnt[nums[i-1]] += 1 + cnt[nums[i-1]-1]
            left[i] = cnt[nums[i]-1]
        }

        cnt = make(map[int]int64)

        for i := n - 2; i >= 0; i-- {
            cnt[nums[i+1]] += 1 + cnt[nums[i+1]+1]
            right[i] = cnt[nums[i]+1]
        }

        var ans int64
        for i, x := range nums {
            ans = (ans + (left[i]+right[i]+(left[i]*right[i]%mod))*int64(x)%mod) % mod
        }
        return ans
    }

    x := calc(nums)
    for i, j := 0, len(nums)-1; i < j; i, j = i+1, j-1 {
        nums[i], nums[j] = nums[j], nums[i]
    }
    y := calc(nums)
    s := int64(0)
    for _, num := range nums {
        s += int64(num)
    }
    return int((x + y + s) % mod)
}

评论