小B的极差之和（青训营X豆包MarsCode）｜豆包MarsCode AI 刷题小B的极差之和问题描述小B拥有一

小B的极差之和

问题描述

小B拥有一个数组 a，她使用这个数组构造了一个新数组 b。其中，a[i] 表示在新数组 b 中有 a[i] 个 i+1。例如，若 a = [2, 3, 1]，那么新数组 b = [1, 1, 2, 2, 2, 3]，因为 a[0] = 2 代表数组 b 中有 2 个 1，a[1] = 3 代表数组 b 中有 3 个 2，a[2] = 1 代表数组 b 中有 1 个 3。

现在，你需要帮助小B求出 b 数组中所有连续子数组的极差之和。由于答案可能非常大，请对 $10^9 + 7$ 取模。

数组的极差定义为子数组的最大值减去最小值。

测试样例

样例1：输入：n = 2,a = [2, 1] 输出：2

样例2：输入：n = 3,a = [1, 2, 1] 输出：6

样例3：输入：n = 4,a = [2, 3, 1, 1] 输出：26

问题分析

我们需要求出数组 b 的所有连续子数组的极差之和，极差定义为子数组中的最大值减去最小值。

给定数组 a，它描述了如何构造数组 b，数组 b 中有 a[i] 个 i+1。例如，如果 a = [2, 3, 1]，那么数组 b 将是 [1, 1, 2, 2, 2, 3]。

目标：

构造数组 b。
计算所有连续子数组的极差之和。

由于构造数组 b 可能非常大（最大可能的长度为 sum(a)），直接遍历所有子数组并计算其极差是不现实的。我们需要通过一种优化的方法来处理这个问题。

解决方案

构造数组 b：直接根据数组 a 构造数组 b，每个元素 i+1 出现 a[i] 次。
计算极差：我们需要考虑如何快速计算数组 b 中所有子数组的极差之和。

- 对于每个元素 x，我们需要知道它作为子数组最大值或最小值的贡献。
- 假设元素 x 在某些子数组中是最大值或最小值，那么我们需要计算这个元素作为最大值或最小值的子数组的数量。

优化计算：可以通过单调栈来计算每个元素作为最大值和最小值出现的次数，进而计算其对极差的贡献。

步骤

构造数组 b：通过遍历 a，生成数组 b。
计算贡献：

- 最大值贡献：使用单调栈计算每个元素作为最大值出现的次数。
- 最小值贡献：使用单调栈计算每个元素作为最小值出现的次数。

计算极差之和：

- 对于每个元素，计算其作为最大值和最小值的贡献。
- 计算极差之和时，按照每个元素的最大值贡献和最小值贡献分别加权。

模取：由于结果可能非常大，需要对每个计算结果取模 10^9 + 7。

代码实现

MOD = 10**9 + 7

def solution(n, a):
    # Step 1: 构建数组b
    b = []
    for i in range(n):
        b.extend([i + 1] * a[i])

    m = len(b)
    total_diff_sum = 0
    
    # Step 2: 计算每个元素在某个子数组中的贡献最大值
    def calculate_contribution(is_max):
        # 将每个元素的贡献计算为max（如果is_max=True）或min（如果is-max=False）
        stack = []
        contribution = [0] * m
        
        for i in range(m):
            while stack and (b[stack[-1]] < b[i] if is_max else b[stack[-1]] > b[i]):
                stack.pop()
            
            if stack:
                contribution[i] = (i - stack[-1]) * (m - i)
            else:
                contribution[i] = (i + 1) * (m - i)
            
            stack.append(i)
        
        return contribution

    # Step 3: 将总贡献计算为最大值和最小值
    max_contrib = calculate_contribution(True)
    min_contrib = calculate_contribution(False)
    
    # Step 4: 计算总差值和
    for i in range(m):
        total_diff_sum += max_contrib[i] - min_contrib[i]
        total_diff_sum %= MOD
    
    return total_diff_sum

解释

构造数组 b：

- 我们通过遍历数组 a，根据每个值 a[i] 构造数组 b，例如如果 a[i] = 2，那么数组 b 中将有两个 i+1。

计算最大值贡献和最小值贡献：

- 使用单调栈来分别计算每个元素作为最大值和最小值的贡献次数。
- 对于每个元素，我们计算它在作为最大值或最小值时，出现在子数组中的次数。

计算极差之和：

- 对于每个元素，计算其作为最大值和最小值的贡献，然后累加其贡献差。

结果模取：为了防止结果溢出，我们在每次累加时都进行模运算。

复杂度分析

时间复杂度：

- 构造数组 b 的时间复杂度是 O(n)，其中 n 是数组 a 的长度。
- 计算最大值和最小值的贡献都使用了单调栈，时间复杂度为 O(m)，其中 m 是数组 b 的长度。
- 总的时间复杂度是 O(n + m)，其中 m = sum(a)，即数组 b 的长度。

空间复杂度：

- 需要存储数组 b 和贡献数组，空间复杂度为 O(m)。

通过单调栈的优化，避免了暴力枚举所有子数组，从而有效减少了计算复杂度。

小B的极差之和（青训营X豆包MarsCode） ｜ 豆包MarsCode AI 刷题