Problem statement

https://binarysearch.com/problems/Lego-Towers/

Solution

Sort heights and then for each window of size k check how much we need to pay.

Complexity

It is O(n log n) for time and O(n) for space.

Code

class Solution:
    def solve(self, H, k):
        H = sorted(H)
        n = len(H)
        acc = [0] + list(accumulate(H))
        ans = float("inf")
        for i in range(n - k + 1):
            ans = min(ans, k*H[i+k-1] - (acc[i + k] - acc[i]))
        
        return ans