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

173 阅读5分钟

问题分析

www.marscode.cn/practice/vk…

  1. 题意

    • 给定一个数组 a,通过该数组构造新数组 b。数组 a[i] 表示数组 b 中有 a[i] 个元素 i+1。然后,我们需要计算数组 b 中所有连续子数组的极差之和(即最大值与最小值之差的总和)。
    • 极差的计算方式:对于每个子数组,极差是该子数组中的最大值减去最小值。计算所有子数组的极差并求和。
  2. 直接暴力解法的难点

    • 如果直接通过暴力枚举数组 b 的所有子数组,并计算每个子数组的极差,这将是一个非常低效的解决方案。假设 b 数组的长度是 m,那么所有子数组的数量为 O(m^2),每个子数组的极差计算需要遍历一遍,这就导致总时间复杂度为 O(m^3),对大输入规模来说不现实。
    • 因此,我们需要一种更高效的方法来计算极差之和,避免显式地遍历所有子数组。

解题思路

为了解决这一问题,我们采用了 单调栈 的技巧,并结合区间贡献法来高效计算每个元素在所有子数组中作为最大值和最小值的贡献。

1. 构造数组 b

首先,我们需要根据数组 a 构造数组 b。每个 a[i] 表示数字 i+1b 中出现的次数。因此,我们可以通过将 i+1 重复 a[i] 次来构建数组 b

  • 例如,如果 a = [2, 3, 1],那么 b = [1, 1, 2, 2, 2, 3]

构造过程显而易见,时间复杂度为 O(n)

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

接下来的关键任务是计算每个元素在所有连续子数组中的极差贡献。为了做到这一点,我们通过 单调栈 来计算每个元素在 b 中作为最大值和最小值的贡献。

  • 最大值贡献: 对于每个元素 b[i],我们需要知道它在多少个子数组中是最大值。可以利用单调栈来找到每个元素左边第一个比它大的元素和右边第一个比它大的元素。

    具体来说:

    • left_max[i]:表示 b[i] 左边第一个大于它的元素的索引。
    • right_max[i]:表示 b[i] 右边第一个大于它的元素的索引。

    每个元素 b[i] 的最大值贡献是它作为最大值出现在的所有子数组的总和,贡献值是 b[i] * (i - left_max[i]) * (right_max[i] - i),即它作为最大值出现在的区间数。

  • 最小值贡献: 类似地,我们可以通过单调栈计算每个元素作为最小值的贡献:

    • left_min[i]:表示 b[i] 左边第一个小于它的元素的索引。
    • right_min[i]:表示 b[i] 右边第一个小于它的元素的索引。

    每个元素 b[i] 的最小值贡献是它作为最小值出现在的所有子数组的总和,贡献值是 b[i] * (i - left_min[i]) * (right_min[i] - i)

3. 总极差之和的计算

最终的极差之和是所有最大值贡献与最小值贡献之差。我们通过遍历 b 中的每个元素,计算其作为最大值和最小值的贡献,最后相加得到总的极差之和。

4. 处理大数问题

由于题目要求结果对 10^9 + 7 取模,所以在计算每个贡献时,我们要确保每个步骤都进行取模运算,防止结果溢出。

算法设计

  • 构造数组 b:根据数组 a 构造 b,时间复杂度 O(n)
  • 单调栈计算最大值和最小值贡献:两次单调栈的遍历计算时间复杂度为 O(m),其中 m 是数组 b 的长度。
  • 总极差之和的计算:遍历 b,计算每个元素的贡献,时间复杂度 O(m)

综上所述,整个算法的时间复杂度为 O(m),其中 m 是构造数组 b 后的长度。由于 m >= n,因此整体时间复杂度是 O(n + m),相对高效。

关键技巧

  1. 单调栈:利用单调栈找出每个元素的左右边界,从而高效计算其在所有子数组中的贡献。
  2. 区间贡献法:通过计算每个元素作为最大值和最小值的贡献,避免了暴力枚举所有子数组,显著提高了效率。

代码实现

MOD = 10**9 + 7

def solution(n: int, a: list) -> int:
    # Step 1: 构造数组 b
    b = []
    for i in range(n):
        b.extend([i + 1] * a[i])
    
    # Step 2: 计算每个元素作为最大值和最小值的贡献
    def calculate_contributions(b):
        n = len(b)
        left_max = [-1] * n  # left_max[i]: the index of the nearest element larger than b[i] on the left
        right_max = [n] * n  # right_max[i]: the index of the nearest element larger than b[i] on the right
        left_min = [-1] * n  # left_min[i]: the index of the nearest element smaller than b[i] on the left
        right_min = [n] * n  # right_min[i]: the index of the nearest element smaller than b[i] on the right
        
        # 单调栈计算
        stack = []
        for i in range(n):
            while stack and b[stack[-1]] <= b[i]:
                stack.pop()
            if stack:
                left_max[i] = stack[-1]
            stack.append(i)
        
        stack.clear()
        for i in range(n - 1, -1, -1):
            while stack and b[stack[-1]] < b[i]:
                stack.pop()
            if stack:
                right_max[i] = stack[-1]
            stack.append(i)
        
        stack.clear()
        for i in range(n):
            while stack and b[stack[-1]] >= b[i]:
                stack.pop()
            if stack:
                left_min[i] = stack[-1]
            stack.append(i)
        
        stack.clear()
        for i in range(n - 1, -1, -1):
            while stack and b[stack[-1]] > b[i]:
                stack.pop()
            if stack:
                right_min[i] = stack[-1]
            stack.append(i)
        
        # 计算每个元素作为最大值和最小值的贡献
        max_contrib = 0
        min_contrib = 0
        
        for i in range(n):
            max_contrib += (i - left_max[i]) * (right_max[i] - i) % MOD * b[i] % MOD
            max_contrib %= MOD
            min_contrib += (i - left_min[i]) * (right_min[i] - i) % MOD * b[i] % MOD
            min_contrib %= MOD
        
        return max_contrib - min_contrib
    
    return calculate_contributions(b) % MOD