LeetCode 2551. Put Marbles in Bags

January 29, 2023

周賽330。只想得到dp解，沒想到是腦筋急轉彎。

題目

你有k個背包。輸入一個整數陣列weights，其中weights[i]代表第i顆彈珠的重量。另外還有整數k。

你必須按照以下規則，將彈珠分裝到k個背包中：

背包不可以為空
如果第i和第j顆彈珠都在同一個背包中，則介於i和j的所有彈珠都必須在同一個背包中
如果一個背包裝了i~j的彈珠，則成本為weights[i] + weights[j]

一個分配方案的分數為個背包的成本總和。

求所有方案中最大分數和最小分數的差值。

解法

看到子陣列切割馬上想到dp(i,k)代表從i點切出k個背包的解法，但是weights長度N和k都高達10^5，怎麼看都不可能過。

看看例題2：

weights = [1, 3], k = 2
切成兩個部分[1]和[3]
成本分別為1+1和3+3

分成k個背包，中間有k-1個分割點。對於每個分割點來說，其左邊的元素會作為左方背包的右端點，而右邊的元素會做為右方背包的左端點。然後第一個元素必定是第一個背包的左端點、最後一個元素必定是最後一個背包的右端點。

weights = [1, 3], k = 2
唯一切割點介於[1,3]之間。成本為1+3
1作為第一個背包左端，3作為最後一個背包右邊。成本再加上1+3
分數為4+4=8

因此我們只要窮舉出所有的切割點，分別找出成本最大/最小的k-1個，便可求得分數。
然而第一個和最後一個元素一定會包含在最大/最小分數的切割方案中，必定相減為0，所以不用特別處理。

時間複雜度O(N log N)。空間複雜度O(N)。

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        cut=[a+b for a,b in pairwise(weights)]
        cut.sort(reverse=True)
        
        ans=0
        for i in range(k-1):
            ans+=cut[i]-cut[-1-i]
            
        return ans

也可以用heap來做，在k比較小的情況下可以有效降低複雜度。
維護min heap保存最大的成本，max heap保存最小的成本，最後求和相減。

時間複雜度O(N log k)。空間複雜度O(k)。

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        N=len(weights)
        min_heap=[]
        max_heap=[]
        for i in range(1,N):
            cost=weights[i]+weights[i-1]
            heappush(min_heap,cost)
            heappush(max_heap,-cost)
            if len(min_heap)>=k:
                heappop(min_heap)
                heappop(max_heap) 
                
        return sum(min_heap)+sum(max_heap)

mocowcow

題目

解法