小B的极差之和| 豆包MarsCode AI刷题

一、思路分析

1. 构造数组 b：

   • 首先，根据题目给定的规则，通过遍历数组 a，将对应数量的 i + 1 元素添加到数组 b 中，从而构造出数组 b。例如，当 a = [2, 3, 1]，就会按照规则依次在 b 中添加元素，得到 b = [1, 1, 2, 2, 2, 3]。

2. 计算极差之和的思路：

   • 为了求出数组 b 中所有连续子数组的极差之和，我们采用了一种巧妙的方法，即分别计算每个元素作为子数组中的最大值和最小值时对极差之和的贡献，然后将这些贡献累加起来。
   • 对于每个元素作为最大值的贡献，我们通过维护一个单调栈来实现。当遍历到一个新元素时，如果栈顶元素对应的 b 中的值小于当前元素的值，就说明栈顶元素在当前位置之前找到了一个更大的元素，此时可以根据栈顶元素的位置以及当前元素的位置计算出栈顶元素作为最大值时对极差之和的贡献，然后将栈顶元素弹出。遍历完整个数组后，还需要处理栈中剩余的元素，同样按照相应规则计算它们作为最大值时的贡献。
   • 同理，对于每个元素作为最小值的贡献，也是通过类似的单调栈操作来实现，只是在比较栈顶元素和当前元素时，是判断栈顶元素对应的 b 中的值是否大于当前元素的值。
   • 最后，将每个元素作为最大值和最小值的贡献之差乘以该元素的值，再累加到总的极差之和 total_sum 中，并对 10**9 + 7 取模，得到最终结果。

二、图解（以 n = 4, a = [2, 3, 1, 1] 为例）

1. 构造数组 b 阶段：

• 按照规则，a[0] = 2，所以在 b 中添加 2 个 1，此时 b = [1, 1]。

• a[1] = 3，在 b 中添加 3 个 2，得到 b = [1, 1, 2, 2, 2]。

• a[2] = 1，添加 1 个 3，变为 b = [1, 1, 2, 2, 2, 3]。

• a[3] = 1，再添加 1 个 4，最终 b = [1, 1, 2, 2, 2, 3, 4]。

2. 计算每个元素作为最大值和最小值的贡献阶段：

计算最大值贡献（以部分元素为例） ：

• 假设我们从左到右遍历 b 来计算每个元素作为最大值的贡献。

• 当遍历到 b[2] = 2（此时 i = 2）时，假设栈中已有 b[0] = 1（假设栈顶元素对应的索引为 0），因为 b[0] < b[2]，所以要计算 b[0] 作为最大值的贡献。此时 idx = 0，栈顶弹出后，idx - (stack[-1] if stack else -1) 这里 stack[-1] 为空（因为栈顶已弹出），所以按规则计算为 0 - (-1) = 1，i - idx = 2 - 0 = 2，那么 b[0] 作为最大值的贡献为 1 * 2 = 2。

计算最小值贡献（以部分元素为例） ：

• 同样从左到右遍历 b 来计算每个元素作为最小值的贡献。
• 当遍历到 b[4] = 2（此时 i = 4）时，假设栈中已有 b[3] = 2（假设栈顶元素对应的索引为 3），因为 b[3] > b[4]，所以要计算 b[4] 作为最小值的贡献。此时 idx = 4，栈顶弹出后，idx - (stack[-1] if stack else -1) 这里 stack[-1] 变为之前栈中的下一个元素（假设为 2），那么 idx - (stack[-1] if stack else -1) 为 4 - 2 = 2，i - idx = 4 - 4 = 0，所以 b[4] 作为最小值的贡献为 2 * 0 = 0（这里只是示例，实际完整计算会考虑更多情况）。

三、代码详解

def solution(n: int, a: list) -> int:
    MOD = 10**9 + 7

    # 构造数组b
    b = []
    for i in range(n):
        b.extend([i + 1] * a[i])

    total_sum = 0
    length = len(b)

    # 计算每个元素作为最大值和最小值的贡献
    def calculate_contribution():
        nonlocal total_sum
        # 计算每个元素作为最大值的贡献
        max_contrib = [0] * length
        min_contrib = [0] * length

        # 计算最大值贡献
        stack = []
        for i in range(length):
            while stack and b[stack[-1]] < b[i]:
                idx = stack.pop()
                max_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (i - idx)
            stack.append(i)

        while stack:
            idx = stack.pop()
            max_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (length - idx)

        # 计算最小值贡献
        stack = []
        for i in range(length):
            while stack and b[stack[-1]] > b[i]:
                idx = stack.pop()
                min_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (i - idx)
            stack.append(i)

        while stack:
            idx = (stack.pop())
            min_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (length - idx)

        # 计算极差之和
        for i in range(length):
            total_sum += (max_contrib[i] - min_contrib[i]) * b[i]
            total_sum %= MOD

    calculate_contribution()

    return total_sum

1. MOD = 10**9 + 7：

   • 定义了一个常量 MOD，用于后续对计算结果取模，以避免结果过大导致数据溢出等问题。

2. b = [] 及后续构造 b 的循环：

   • 首先创建一个空列表 b，然后通过遍历数组 a，对于每个 a[i]，使用 extend 方法将 [i + 1] 重复 a[i] 次添加到 b 中，从而按照题目规则构造出数组 b。

3. calculate_contribution 函数：

   • 这是核心函数，用于计算每个元素作为最大值和最小值的贡献，并最终计算出极差之和。

   • max_contrib = [0] * length 和 min_contrib = [0] * length：创建两个长度为 b 数组长度的列表，分别用于存储每个元素作为最大值和最小值的贡献值，初始化为 0。

   • 计算最大值贡献部分：

     • stack = []：创建一个空栈，用于辅助计算每个元素作为最大值的贡献。
     • for i in range(length) 循环：遍历数组 b。在循环中，当 stack 不为空且 b[stack[-1]] < b[i] 时，说明栈顶元素对应的 b 中的值小于当前元素的值，此时需要计算栈顶元素作为最大值的贡献。通过 idx = stack.pop() 弹出栈顶元素，然后按照公式 max_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (i - idx) 计算其贡献值，并更新 max_contrib 列表。最后将当前元素的索引 i 压入栈中。循环结束后，还需要处理栈中剩余的元素，同样按照公式 max_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (length - idx) 计算它们作为最大值的贡献值。

   • 计算最小值贡献部分：

     • 与计算最大值贡献部分类似，也是通过一个空栈和类似的循环来实现。只是在比较条件上，是当 stack 不为空且 b[stack[-1]] > b[i] 时，计算栈顶元素作为最小值的贡献，公式为 min_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (i - idx)，循环结束后处理栈中剩余元素的公式为 min_contrib[idx] = (idx - (stack[-1] if stack else -1)) * (length - idx)。

   • 计算极差之和部分：

     • 通过 for i in range(length) 循环遍历数组 b，对于每个元素 i，将其作为最大值和最小值的贡献之差乘以该元素的值，即 (max_contrib[i] - min_contrib[i]) * b[i]，累加到 total_sum 中，并对 MOD 取模，以保证结果在规定范围内。

四、知识总结

新知识点梳理与分析

1. 根据特定规则构造数组：

   • 本题中根据数组 a 按照特定规则构造数组 b 的过程，涉及到对数组的遍历以及根据元素值重复添加特定元素到新数组的操作。这种根据给定规则生成新数组的方式在很多编程问题中都有应用，它要求我们能够准确理解规则并将其转化为代码实现。例如，在处理一些数据转换、数据生成相关的任务时，可能会遇到类似需要根据已有数据按照特定模式生成新数据的情况。

2. 单调栈的应用：

   • 在计算每个元素作为最大值和最小值的贡献时，使用了单调栈这一数据结构。单调栈在处理与数组元素的相对大小关系以及计算元素在特定区间内的贡献等问题上非常有效。它通过维护栈内元素的单调性质（本题中分别维护了单调递增和单调递减的栈用于计算最大值和最小值贡献），可以快速判断元素之间的相对大小关系，并根据栈的状态和当前元素的位置计算出元素的贡献值。单调栈的应用场景较为广泛，比如在计算数组中的下一个更大元素、前一个更大元素等问题中都可以发挥作用。

理解与学习建议

• 理解：

  • 本题的关键在于理解如何根据给定规则构造数组 b，以及如何利用单调栈来高效地计算每个元素作为最大值和最小值的贡献，进而求出所有连续子数组的极差之和。构造数组 b 的过程相对直观，只要按照规则依次处理数组 a 中的元素即可。而利用单调栈计算贡献的部分则需要深入理解单调栈的工作原理，即通过栈内元素的单调性质来快速判断元素的相对大小关系，并根据相应公式计算贡献值。

• 学习建议：

  • 对于初学者来说，首先要熟练掌握数组的基本操作，如遍历、添加元素、根据索引获取元素等。这些基本操作是实现本题代码的基础，例如在构造数组 b 时就大量用到了数组的遍历和添加元素操作。
  • 深入学习单调栈的相关知识。了解单调栈的概念、基本原理以及常见的应用场景。可以通过多做一些涉及单调栈应用的练习题，如计算数组中的下一个更大元素、前一个更小元素等问题，来加深对单调栈的理解和运用能力。在理解本题中利用单调栈计算贡献的代码时，可以自己手动模拟一下计算过程，比如按照代码中的步骤，在纸上画出栈的状态变化以及元素贡献值的计算过程，这样可以更直观地理解代码的运行机制。
  • 在遇到类似需要根据特定规则生成新数组或利用特定数据结构解决问题的题目时，不要急于动手写代码，先仔细分析题目规则和要求，思考有哪些数据结构或算法可能适合解决该问题。然后再逐步将思路转化为代码实现，这样可以提高代码的编写效率和质量。