在K个连续子数组的最小值中最大化最大值

问题描述

给定一个长度为 N 的整数序列和一个正整数 K，你需要将这个序列划分成 K 个连续的子数组，使得子数组中的最小值的和最大。例如，对于序列 [1, 3, 1, 2, 2, 2]，当 K=3 时最优的划分为 [1, 3], [1, 2], [2, 2]，这样子数组中的最小值之和为 1+1+2=4。

解法

思路

我们可以使用二分答案来解决这个问题。我们假设当前的答案为 X，然后我们就需要找到一种划分方法，使得这种划分方法中的子数组最小值的和不大于 X。

那么我们可以采用贪心的思想来解决这个问题：尽可能地将连续的元素合并成一个子数组，同时保证这个子数组的最小值不大于 X。如果我们找到的划分方案中连续的子数组的数量小于 K，那么说明我们的答案 X 过大了；如果连续的子数组的数量大于等于 K，说明我们的答案 X 可以更小一些。

为了使连续的元素能够很方便地合并成一个子数组，我们可以使用单调栈来统计前缀数组中的每个元素的最左边的小于它的元素和最右边的小于它的元素。然后对于每个元素，我们就可以将它作为一个分界点，分成左边和右边两个子数组，使得每个子数组的值都不大于 X。最后我们将左右两个子数组的贡献合并起来，就得到了当前方案的最小值之和。

代码

from typing import List


def check(nums: List[int], k: int, x: int) -> bool:
    """判断是否存在一种划分方法，使得子数组中的最小值之和不大于 x"""
    stack = []
    left = [-1] * len(nums)
    for i in range(len(nums)):
        while stack and nums[stack[-1]] >= nums[i]:
            stack.pop()
        if stack:
            left[i] = stack[-1]
        stack.append(i)

    stack = []
    right = [len(nums)] * len(nums)
    for i in range(len(nums) - 1, -1, -1):
        while stack and nums[stack[-1]] >= nums[i]:
            stack.pop()
        if stack:
            right[i] = stack[-1]
        stack.append(i)

    ans = 0
    i = 0
    while i < len(nums) and k > 0:
        j = i
        while j < len(nums) and left[j] >= i and right[j] <= i + k - 1:
            j += 1
        ans += (j - i) * nums[i]
        i = j
        k -= 1
    return ans <= x * (len(nums) // k)


def max_subarray_min_sum(nums: List[int], k: int) -> int:
    left, right = min(nums), sum(nums)
    while left <= right:
        mid = (left + right) // 2
        if check(nums, k, mid):
            left = mid + 1
        else:
            right = mid - 1
    return left - 1

代码中的 check 函数用来判断是否存在一种划分方法，使得子数组中的最小值之和不大于 x；max_subarray_min_sum 函数用来求解最小的 x，使得满足所有的划分方案中子数组中的最小值的和不大于 x。

时间复杂度

我们需要进行 O(log sum(nums)) 次二分查找，每次二分查找需要进行 O(len(nums)) 次 check 操作，而 check 操作的时间复杂度为 O(len(nums))，因此，总的时间复杂度为 O(len(nums) * log sum(nums))。