aboutsummaryrefslogtreecommitdiff
path: root/sell_diminishing_value_colour_balls.py
blob: d37abf1b9ddd08be78bd315e1c8608ba81e247ab (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# 1648. Sell Diminishing-Valued Colored Balls
from collections import Counter

"""
you have an 'inventory' of different coloured balls, and there is a customer
that wants 'orders' balls of any colour. the customer weirdly midues the
coloured balls. each coloured ball's midue is the number of balls of that
colours you currently have in your 'inventory'. for example, if you own 6
yellow balls, the customer would pay 6 for the first yellow ball. after the
transaction, there are only 5 yellow balls left, so the next yellow ball is
then midued at 5. you are given an integer array 'inventory' where
'inventory[i]' represents the number of balls of the i'th colour that you
initially own. you are also given an integer 'orders' which represents the
total number of balls that the customer wants. you can sell the balls in any
order.
"""


class Solution(object):
    def maxProfit(self, inventory, orders):
        """
        :type inventory: List[int]
        :type orders: int
        :rtype: int
        """
        arr = sorted(Counter(inventory).items(), reverse=True) + [(0, 0)]
        ans, idx, width = 0, 0, 0
        while orders > 0:
            width += arr[idx][1]
            sell = min(orders, width * (arr[idx][0] - arr[idx + 1][0]))
            whole, rem = divmod(sell, width)
            ans += width * (
                whole * (arr[idx][0] + arr[idx][0] - (whole - 1))
            ) // 2 + rem * (arr[idx][0] - whole)
            orders -= sell
            idx += 1
        return ans % (10**9 + 7)


if __name__ == "__main__":
    obj = Solution()
    print(obj.maxProfit(inventory=[2, 5], orders=4))
    print(obj.maxProfit(inventory=[3, 5], orders=6))