23.07.08 Today’s Leetcode

2551. Put Marbles in Bags (hard)

TLE at 96/106

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        n = len(weights)

        @functools.lru_cache(None)
        def dfs(index, count):
            if count > k:
                return [math.inf, -math.inf]
            if index >= n:
                return [math.inf, -math.inf]
            if count == k - 1:
                p = weights[index] + weights[n - 1]
                return [p, p]
            
            temp = [math.inf, -math.inf]
            for i in range(index, n):
                num = weights[index] + weights[i]
                min_, max_ = dfs(i + 1, count + 1)
                temp[0] = min(temp[0], num + min_)
                temp[1] = max(temp[1], num + max_)
            
            return temp

        res = dfs(0, 0)
        # print(dfs(0, 0))

        return res[1] - res[0]

fucking amazing solution

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        arr = sorted(map(sum, (pairwise(weights))))
        return  sum(arr[-k+1:]) - sum(arr[:k-1])