Problem statement

https://leetcode.com/problems/sell-diminishing-valued-colored-balls/

Solution

The idea is to answer question: given value mid: how many balls with value >= mid we can take. If we draw numbers of balls with bars, then there will be the place, when we have fully taken rows on the top and then there will be one not fully filled row. We can find the place of the last fully filled row using binary search.

Complexity

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

Code

class Solution:
    def maxProfit(self, I, order):
        beg, end = 0, max(I) + 1
        
        while beg + 1 < end:
            mid = (beg + end)//2
            if sum(max(i-mid, 0) for i in I) <= order:
                end = mid
            else:
                beg = mid
                
        ans = sum(max(0, (i + end + 1)*(i - end)) for i in I)//2
        rest = order - sum(max(i - end, 0) for i in I)   
        return (ans + rest * end) % (10**9 + 7)