2025-08-02：最多 K 个元素的子数组的最值之和。用go语言，给定一个整数数组 nums 和一个正整数 k，请找出所有长

福大大架构师每日一题

发布于 2025-08-05 15:12:43

24000

代码可运行

文章被收录于专栏：福大大架构师每日一题福大大架构师每日一题

运行总次数：0

代码可运行

2025-08-02：最多 K 个元素的子数组的最值之和。用go语言，给定一个整数数组 nums 和一个正整数 k，请找出所有长度最多为 k 的连续子数组，计算这些子数组中最大值和最小值的和，并返回最大的那个和。

1 <= nums.length <= 80000。

1 <= k <= nums.length。

-1000000 <= nums[i] <= 1000000。

输入：nums = [1,2,3], k = 2。

输出：20。

解释：

最多 2 个元素的 nums 的子数组：

当然，下面是你提供的数据转换成的表格形式：

子数组	最小值	最大值	和
[1]	1	1	2
[2]	2	2	4
[3]	3	3	6
[1, 2]	1	2	3
[2, 3]	2	3	5
总和			20

输出为 20 。

题目来自力扣3430。

解决思路

直接暴力计算所有子数组的最大值和最小值的和显然是不现实的，因为子数组的数量是 O(n * k)，对于 n = 80000 和 k = 80000 来说，这会非常慢。因此，我们需要一种更高效的方法。

关键观察

1. 单调栈：我们可以利用单调栈来计算所有子数组的最小值的贡献和最大值的贡献。具体来说：
- • 对于最小值，我们需要计算每个元素作为最小值出现在多少个子数组中，然后乘以该元素的值，得到其贡献。
- • 对于最大值，类似地，我们需要计算每个元素作为最大值出现在多少个子数组中，然后乘以该元素的值，得到其贡献。
- • 最终结果是所有子数组的最大值的和加上所有子数组的最小值的和（即 sum(max) + sum(min)）。
2. 限制子数组长度：我们需要限制子数组的长度不超过 k。因此，在计算某个元素作为最小值或最大值时，需要考虑子数组的长度限制。

具体步骤

1. 计算最小值的贡献：
- • 使用单调递增栈来找到每个元素作为最小值可以覆盖的范围（即左右边界）。
- • 对于每个元素 nums[i]，找到左边第一个比它小的元素的位置 l 和右边第一个比它小的元素的位置 r。这样，nums[i] 是 [l+1, r-1] 区间内的最小值。
- • 计算以 nums[i] 为最小值的子数组的数量（考虑长度不超过 k），然后乘以 nums[i] 得到其贡献。
- • 具体来说，子数组的数量可以通过组合数学计算：在 [l+1, r-1] 区间内，包含 nums[i] 的子数组的数量，且长度不超过 k。
2. 计算最大值的贡献：
- • 类似地，我们可以通过将所有元素取反，然后复用计算最小值的逻辑来计算最大值的贡献。
- • 因为 max(a, b) = -min(-a, -b)，所以对 nums 取反后计算最小值等价于原数组的最大值。
3. 合并结果：
- • 计算所有子数组的最小值的和 sum_min。
- • 计算所有子数组的最大值的和 sum_max（通过取反后计算 sum_min 得到）。
- • 最终结果是 sum_max + sum_min。

子数组数量的计算

对于一个元素 nums[i]，其作为最小值的范围是 [l+1, r-1]，其中 l 是左边第一个比它小的位置，r 是右边第一个比它小的位置。子数组的数量可以通过以下方式计算：

• 子数组必须包含 nums[i]，且长度不超过 k。
• 设 m = r - l - 1 是 nums[i] 作为最小值的区间长度。
- • 如果 m <= k，则子数组数量为 m * (m + 1) / 2（即所有可能的子数组）。
- • 如果 m > k，则需要考虑长度限制。子数组的数量为 (2*m - k + 1) * k / 2。

时间复杂度

• 单调栈的处理是 O(n) 的，因为每个元素最多入栈和出栈一次。
• 对于每个元素，计算子数组数量的操作是 O(1) 的。
• 因此，总的时间复杂度是 O(n)。

额外空间复杂度

• 单调栈的空间是 O(n)。
• 其他临时变量的空间是 O(1)。
• 因此，总的额外空间复杂度是 O(n)。

总结

1. 使用单调栈计算每个元素作为最小值和最大值的贡献。
2. 通过数学公式计算限制子数组长度不超过 k 时的子数组数量。
3. 合并最小值和最大值的贡献得到最终结果。
4. 时间复杂度：O(n)。
5. 额外空间复杂度：O(n)。

Go完整代码如下：

package main

import (
    "fmt"
    "math"
)

func minMaxSubarraySum(nums []int, k int)int64 {
    count := func(m int)int {
        if m <= k {
            return (m + 1) * m / 2
        }
        return (m*2 - k + 1) * k / 2
    }

    // 计算最小值的贡献
    sumSubarrayMins := func() (res int) {
        st := []int{-1} // 哨兵
        for r, x := range nums {
            forlen(st) > 1 && nums[st[len(st)-1]] >= x {
                i := st[len(st)-1]
                st = st[:len(st)-1]
                l := st[len(st)-1]
                cnt := count(r-l-1) - count(i-l-1) - count(r-i-1)
                res += nums[i] * cnt // 累加贡献
            }
            st = append(st, r)
        }
        return
    }

    nums = append(nums, math.MinInt)
    ans := sumSubarrayMins()
    // 所有元素取反（求最大值），就可以复用同一份代码了
    for i := range nums {
        nums[i] = -nums[i]
    }
    ans -= sumSubarrayMins()
    returnint64(ans)
}

func main() {
    nums := []int{1, 2, 3}
    k := 2
    result := minMaxSubarraySum(nums, k)
    fmt.Println(result)
}

Python完整代码如下：

# -*-coding:utf-8-*-

import math

def minMaxSubarraySum(nums, k):
    def count(m):
        if m <= k:
            return (m + 1) * m // 2
        return (m * 2 - k + 1) * k // 2

    def sumSubarrayMins():
        res = 0
        st = [-1]  # 哨兵
        for r in range(len(nums)):
            x = nums[r]
            while len(st) > 1 and nums[st[-1]] >= x:
                i = st.pop()
                l = st[-1]
                cnt = count(r - l - 1) - count(i - l - 1) - count(r - i - 1)
                res += nums[i] * cnt  # 累加贡献
            st.append(r)
        return res

    # 计算最小值的贡献
    original_nums = nums.copy()
    nums.append(-math.inf)
    ans = sumSubarrayMins()
    
    # 所有元素取反（求最大值），就可以复用同一份代码了
    nums = original_nums.copy()
    nums = [-x for x in nums]
    nums.append(-math.inf)
    ans -= sumSubarrayMins()
    
    return ans

nums = [1, 2, 3]
k = 2
result = minMaxSubarraySum(nums, k)
print(result)