1648. Sell Diminishing Valued Colored Balls (Medium) - TengnanYao/daily_leetcode GitHub Wiki

class Solution(object):
    def maxProfit(self, inventory, orders):
        """
        :type inventory: List[int]
        :type orders: int
        :rtype: int
        """
        def getSum(l, r):
            return (l + r) * (r - l + 1) // 2
        def countSum(j, base):
            return sum(inventory[i] - base for i in range(j, len(inventory)))
        
        inventory.sort()
        l, r = 0, max(inventory)
        while l <= r:
            m = (l + r) // 2
            i = bisect_right(inventory, m)
            j = bisect_right(inventory, m - 1)
            s1 = countSum(i, m)
            if s1 <= orders <= s1 + len(inventory) - j:
                return (sum(getSum(m + 1, inventory[k]) for k in range(i, len(inventory))) + (orders - s1) * m) % (10 ** 9 + 7)
            elif orders < s1:
                l = m + 1
            else:
                r = m - 1