最多K个相邻元素的最大子序列和

简介

在求最大子序列和的基础上，考虑对不能相邻的元素的个数做出限制，即最多只能有K个连续的元素。

解题思路

一、动态规划

基本思路：用dp[i]表示以第i个元素结尾的最大子序列和，则状态转移方程为：

dp[i] = max(dp[k]) + nums[i] (i-k > K)

其中，dp[k]为以第k个元素结尾的最大子序列和，k满足条件i-k > K。时间复杂度为O(n^2)。

优化： 观察状态转移方程可知，dp[i]只和前K+1个位置的dp值有关系，因此可以用单调队列优化。

单调队列：一个队列，其中的元素按照顺序单调递增或者单调递减。

对于本题，我们维护单调递减的队列，队尾对应着当前子序列中最小的元素。

将队列中的元素下标按照dp值的大小从小到大排序。

每次遍历到i位置时，从队尾开始弹出元素，直到队列中只剩下不满足i-k<=K的元素，即保证了队列中的元素下标之差最大为K。

则dp[i] = dp[队列中dp值最大的元素下标] + nums[i]。

时间复杂度为O(n)。

二、分治法

分治法的核心思想就是将问题分解成规模更小的子问题，然后分别解决子问题，最后将子问题的解合并成原始问题的解。

考虑将数组a分成左右两段，递归求解左半段的最大子序列和、右半段的最大子序列和，以及跨越中点的最大子序列和。

对于后者，我们可以从中点向左、向右分别遍历，求出左、右两侧连续的最大子序列和，然后将它们加起来，即可得到跨越中点的最大子序列和。

最后，再从三者中选择出最大值即为整个数组的最大子序列和。

时间复杂度为O(nlogn)。

代码实现

动态规划

def max_subarray(nums: List[int], K: int) -> int:
    n = len(nums)
    dp = [-float('inf')] * n
    queue = []
    for i in range(n):
        while queue and i-queue[0]>K:
            queue.pop(0)
        if queue:
            j = queue[0]
            dp[i] = dp[j] + nums[i]
        dp[i] = max(dp[i], nums[i])
        while queue and dp[i]<=dp[queue[-1]]:
            queue.pop()
        queue.append(i)
    return max(dp)

分治法

def max_crossing_subarray(a: List[int], left: int, mid: int, right: int) -> int:
    # 左侧最大子序列和
    left_sum = -float('inf')
    cur_sum = 0
    for i in range(mid, left-1, -1):
        cur_sum += a[i]
        left_sum = max(left_sum, cur_sum)
    # 右侧最大子序列和
    right_sum = -float('inf')
    cur_sum = 0
    for j in range(mid+1, right+1):
        cur_sum += a[j]
        right_sum = max(right_sum, cur_sum)
    return left_sum + right_sum

def max_subarray(a: List[int], left: int, right: int, K: int) -> int:
    if left == right:
        return a[left]
    mid = (left + right) // 2
    # 左侧最大子序列和
    left_sum = max_subarray(a, left, mid, K)
    # 右侧最大子序列和
    right_sum = max_subarray(a, mid+1, right, K)
    # 跨越中点的最大子序列和
    cross_sum = max_crossing_subarray(a, left, mid, right)
    return max(left_sum, right_sum, cross_sum)

def max_subarray(nums: List[int], K: int) -> int:
    n = len(nums)
    if K >= n:
        return sum(nums)
    # 将间隔大于K的元素置为0
    for i in range(K+1, n, K+1):
        nums[i] = 0
    return max_subarray(nums, 0, n-1, K)