又是隨便抽到的DP,只有推出狀態轉移,但是不知道怎麼優化,看來要化簡轉移還是有點難度。

題目

輸入整數陣列nums和整數k,回傳nums非空子序列的最大和,且子序列中的每兩個連續整數nums[i]和nums[j],兩者索引距離不得超過k。

解法

先講一開始的TLE版本。
試著以每個數字nums[i],接下來可以選擇num[i+1]~num[i+k]中的最佳解做為下一個數字。

定義dp(i):以nums[i]為起點所能得到的最大和
轉移方程式:dp(i)=nums[i]+max(0,dp(j) FOR ALL i<j<=i+k)
base case:當i為nums中最後一個數時,dp(N-1)=nums[N-1]

但是k跟N最大會到10^5,所以O(N*K)時間是沒辦法過的,必須想辦法化簡。

class Solution:
    def constrainedSubsetSum(self, nums: List[int], k: int) -> int:
        N=len(nums)
        
        @cache
        def dp(i):
            best=0
            for j in range(i+1,i+k+1):
                if j==N:
                    break
                best=max(best,dp(j))
            return nums[i]+best
        
        return max(dp(i) for i in range(N))

看看提示說用queue來優化,多翻了幾篇文章才搞懂意思。
簡單來說就是維護一個單調遞減佇列q,而q的第一個元素永遠是在k距離中的最佳解,如此一來可以降低複雜度至O(N)。

改從尾端進行bottom up,從後方往前對每個數字nums[i]做DP:
每次先將超出k距離的候選人pop掉,然後取最佳解q[0]更新dp[i]的值,將佇列尾端小於dp[i]的值移除,最後將i放入佇列末端。
回傳dp中最大者就是答案。

class Solution:
    def constrainedSubsetSum(self, nums: List[int], k: int) -> int:
        N=len(nums)
        dp=[0]*N
        q=deque()

        for i in range(N-1,-1,-1):
            while q and q[0]-i>k:
                q.popleft()
            best=dp[q[0]] if q else 0
            dp[i]=max(best,0)+nums[i]
            while q and dp[i]>=dp[q[-1]]:
                q.pop()
            q.append(i)
            
        return max(dp)

提示還說了可以用heap優化,我就來試試。
個人感覺是比queue還好寫,但效能確實是差了點,應該是O(N log N),但不知道為什麼都沒有看到用heap的題解,這至少也算個可以接受的答案。

使用max heap保證最佳解一定在頂端,我們只要排除掉超出k距離的候選人就好,不像queue維持單調遞減這麼麻煩。

class Solution:
    def constrainedSubsetSum(self, nums: List[int], k: int) -> int:
        N=len(nums)
        h=[]
        ans=-inf

        for i in range(N-1,-1,-1):
            while h and h[0][1]-i>k:
                heappop(h)
            best=max(0,-h[0][0]) if h else 0
            best+=nums[i]
            ans=max(ans,best)
            heappush(h,[-best,i])   
            
        return ans