子数组的最小值之和

标签: 数组 动态规划 单调栈

难度: Medium

给定一个整数数组 arr,找到 min(b) 的总和,其中 b 的范围为 arr 的每个(连续)子数组。

由于答案可能很大,因此 返回答案模 10^9 + 7

 

示例 1:

输入:arr = [3,1,2,4]
输出:17
解释:
子数组为 [3],[1],[2],[4],[3,1],[1,2],[2,4],[3,1,2],[1,2,4],[3,1,2,4]。 
最小值为 3,1,2,4,1,1,2,1,1,1,和为 17。

示例 2:

输入:arr = [11,81,94,43,3]
输出:444

 

提示:

  • 1 <= arr.length <= 3 * 104
  • 1 <= arr[i] <= 3 * 104

 

Submission

运行时间: 83 ms

内存: 19.2 MB

class Solution:
    def sumSubarrayMins(self, arr: List[int]) -> int:
        mod = 1_000_000_007
        s = [-1]
        arr.append(0)
        ans = 0
        for idx, i in enumerate(arr):
            while i <= arr[s[-1]]:
                t = s.pop()
                if t == -1:
                    return ans
                ans = (ans + arr[t] * (t-s[-1])  * (idx-t)) % mod
            s.append(idx)
        

Explain

本题解使用单调栈的方法来找出所有子数组的最小值之和。首先,将原数组末尾添加一个0,以保证所有元素都能被正确处理。单调栈用于存储数组元素的索引,且栈中的元素对应的数组值将保持单调不增的顺序。对于每个元素,如果它小于或等于栈顶元素,那么栈顶元素及其之前的元素将依次出栈,直到栈顶元素小于当前元素。对于每个出栈的元素t,计算以arr[t]为最小值的子数组的数量,乘以arr[t],再累加到结果中。这样,每个元素在出栈时,可以算出以它为最小值的所有连续子数组对结果的贡献。最终,返回累加的结果,对1,000,000,007取模。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def sumSubarrayMins(self, arr: List[int]) -> int:
        mod = 1_000_000_007  # 定义模数
        s = [-1]  # 初始化栈,用-1作为哨兵元素
        arr.append(0)  # 将0添加到数组末尾,以处理所有元素
        ans = 0  # 初始化结果累加器
        for idx, i in enumerate(arr):  # 遍历数组及其索引
            while i <= arr[s[-1]]:  # 当当前元素小于等于栈顶元素时
                t = s.pop()  # 弹出栈顶元素
                if t == -1:  # 如果栈为空则直接返回结果
                    return ans
                ans = (ans + arr[t] * (t-s[-1])  * (idx-t)) % mod  # 计算以arr[t]为最小值的子数组的贡献,并累加到结果中
            s.append(idx)  # 将当前索引压入栈

Explore

在数组末尾添加一个0的主要作用是确保所有数组中的元素都能被处理并出栈。因为0是一个非常小的值,当遍历到这个0时,它会引起栈中所有当前的元素依次出栈。这样,每个元素都会有机会被计算在内,确保了算法的完整性和正确性。添加0简化了代码逻辑,避免了在循环结束后还需要单独处理栈中剩余元素的复杂情况。

在单调栈中使用-1作为哨兵元素主要是为了处理边界情况,使得栈中始终有一个元素。这样在计算以某个元素为最小值的子数组数量时可以方便地使用索引差来计算。如果栈为空,则需要特别处理,哨兵元素-1简化了这些计算过程,避免了对空栈的额外检查,使得算法更加简洁和高效。

对于每个出栈的元素t,计算以arr[t]为最小值的子数组的数量涉及到两部分:1) t左侧的子数组范围,2) t右侧的子数组范围。具体来说,arr[t]左侧的第一个比arr[t]小的元素的位置是s[-1],因此左侧有(t - s[-1])种可能的起始位置;右侧的第一个比arr[t]小的元素的位置是idx,因此右侧有(idx - t)种可能的结束位置。因此,以arr[t]为最小值的子数组的总数是(t - s[-1]) * (idx - t)。

当栈顶元素等于当前元素时,也让栈顶元素出栈的原理是为了保证所有子数组的最小元素都能被正确计算。如果允许栈里有重复值,则可能会造成重复计算某些子数组的贡献或者遗漏某些情况。通过确保栈中元素严格单调递增,可以简化计算过程,并确保每个子数组的最小值只计算一次,这有助于维护算法的正确性和效率。