[LeetCode in Python] 1425 (H) constrained subsequence sum 带限制的子序列和

题目

https://leetcode-cn.com/problems/constrained-subsequence-sum/

给你一个整数数组 nums 和一个整数 k ，请你返回 非空 子序列元素和的最大值，子序列需要满足：子序列中每两个 相邻 的整数 nums[i] 和 nums[j] ，它们在原数组中的下标 i 和 j 满足 i < j 且 j - i <= k 。
数组的子序列定义为：将数组中的若干个数字删除（可以删除 0 个数字），剩下的数字按照原本的顺序排布。

示例 1：

输入：nums = [10,2,-10,5,20], k = 2
输出：37

解释：子序列为 [10, 2, 5, 20] 。

示例 2：

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

解释：子序列必须是非空的，所以我们选择最大的数字。

示例 3：

输入：nums = [10,-2,-10,-5,20], k = 2
输出：23

解释：子序列为 [10, -2, -5, 20] 。

提示：

1 <= k <= nums.length <= 10^5
-10^4 <= nums[i] <= 10^4

解题思路

• 设dp[i]为以nums[i]结尾的满足要求的子数组最大和
• DP的递推公式是 dp[i] = max(dp[i-1], dp[i-2], ..., dp[i-k], 0) + nums[i]
• 如果直接这么写会TLE，所以需要应用滑动窗口取最大值的方法，O(N)的就是单调队列了
• 单调队列可以参考[LeetCode in Python] 239 (H) sliding window maximum 滑动窗口最大值
• 使用deque来存储k的窗口内dp[]的下标，存下标的好处是不怕相同元素，也方便比较滑窗边缘
• 先对dp和q初始化首元素，可以减少循环中的判断

代码

class Solution:
    def constrainedSubsetSum(self, nums: List[int], k: int) -> int:
        # - dp[i] is max sum of subset and must include nums[i]
        # - formula is: dp[i] = max(dp[i-1], dp[i-2], ..., dp[i-k], 0) + nums[i]
        dp = [0] * len(nums)
        
        # - monotonic queue and keep index of dp[]
        # - q[0] is the index of max dp[]
        q = collections.deque()
        
        # - init dp and q
        dp[0] = nums[0]
        q.append(0)
        
        # - from 1 to n-1
        for i in range(1, len(nums)):
            # - update dp[i]
            dp[i] = max(dp[q[0]], 0) + nums[i]

            # - if window slide, pop the left most index
            if q[0] == i-k: 
                q.popleft()

            # - update q, pop all elements which less than current one
            while q and dp[q[-1]] < dp[i]:
                q.pop()
            q.append(i)

        return max(dp)