Program to find maximum profit by selling diminishing-valued colored balls in Python

Suppose we have an array called, inventory, where inventory[i] represents the number of balls of the ith color we have initially. We also have a value called orders, which represents the total number of balls that the customer wants. we can sell the balls in any order. In our inventory there are different colored balls, customers want balls of any color. Now the values of the balls are special in nature. Each colored ball's value is the number of balls of that color we have in our inventory. So if we currently have 6 blue balls, the customer would pay 6 for the first blue ball. Then there are only 5 blue balls left, so the next blue ball is then valued at 5. We have to find the maximum total value that we can get after selling orders colored balls. If the answer is too large, then, return it modulo 10^9 + 7.

So, if the input is like inventory = [5,7], orders = 6, then the output will be 31 because we can sell the first colored ball twice at price (5,4), and second colored balls 4 times (7,6,5,4), so total profit 5+4+7+6+5+4 = 31

To solve this, we will follow these steps −

• low := 0, high := 10000000

• while low < high, do

  • mid := quotient of (low + high)/2

  • s := 0

  • for each i in inventory, do

    • if i > mid, then

      • s := s + i - mid

  • if s > orders, then

    • low := mid + 1

  • otherwise,

    • high := mid

• mid := quotient of (low + high)/2

• ans := 0

• for each i in inventory, do

  • if i > mid, then

    • ans := ans + quotient of (i*(i+1)/2) - quotient of (mid*(mid+1)/2)

    • orders := orders - i-mid

• ans := ans + orders * mid

• return ans mod (10^9 + 7)

Example

Let us see the following implementation to get better understanding −

def solve(inventory, orders):
   low = 0
   high = 10000000

   while low < high:
      mid = (low+high)//2

      s = 0
      for i in inventory:
         if i > mid:
            s += i-mid

      if s > orders:
         low = mid+1
      else:
         high = mid


   mid = (low+high)//2

   ans = 0
   for i in inventory:
      if i > mid:
         ans += i*(i+1)//2 - mid*(mid+1)//2
         orders -= i-mid

   ans += orders*mid
   return ans % (10**9 + 7)

inventory = [5,7]
orders = 6
print(solve(inventory, orders))

Input

[6,8,7,11,5,9], [0,0,2], [2,3,5]

Output

31