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, and customers want balls of any color.

The values of the balls are special in nature − each colored ball's value equals the current number of balls of that color in our inventory. So if we currently have 6 blue balls, the customer pays 6 for the first blue ball. After selling one, only 5 blue balls remain, so the next blue ball costs 5.

We need to find the maximum total value after selling orders colored balls. If the answer is too large, return it modulo 10^9 + 7.

Example

If the input is inventory = [5,7] and orders = 6, the output will be 31 because we can sell the first colored ball twice at prices (5,4), and second colored balls 4 times at prices (7,6,5,4), giving total profit: 5+4+7+6+5+4 = 31.

Approach Using Binary Search

We use binary search to find the minimum price threshold. The algorithm works by:

• Using binary search to find the lowest price we can sell balls at

• Calculating total profit using arithmetic series formula

• Handling remaining orders at the threshold price

Implementation

def solve(inventory, orders):
    low = 0
    high = 10000000
    
    # Binary search to find minimum price threshold
    while low < high:
        mid = (low + high) // 2
        
        # Count balls we can sell above mid price
        s = 0
        for i in inventory:
            if i > mid:
                s += i - mid
        
        if s > orders:
            low = mid + 1
        else:
            high = mid
    
    # Calculate profit
    mid = (low + high) // 2
    ans = 0
    
    for i in inventory:
        if i > mid:
            # Sum from (mid+1) to i using arithmetic series
            ans += i * (i + 1) // 2 - mid * (mid + 1) // 2
            orders -= i - mid
    
    # Add profit from remaining orders at threshold price
    ans += orders * mid
    return ans % (10**9 + 7)

# Test the function
inventory = [5, 7]
orders = 6
print(solve(inventory, orders))

31

How It Works

The algorithm uses binary search to find the optimal price threshold:

1. Binary Search: Find the minimum price mid such that we can sell exactly orders balls

2. Profit Calculation: For each color with inventory > mid, calculate profit using arithmetic series formula

3. Remaining Orders: Sell any remaining balls at the threshold price

Step-by-Step Example

For inventory = [5,7] and orders = 6:

inventory = [5, 7]
orders = 6

# Binary search finds threshold = 3
# For inventory[0] = 5: sell at prices 5,4 (profit = 9)
# For inventory[1] = 7: sell at prices 7,6,5,4 (profit = 22)
# Total: 9 + 22 = 31

def detailed_example():
    inventory = [5, 7]
    orders = 6
    
    # Simulate the selling process
    total_profit = 0
    balls_sold = 0
    
    print("Selling process:")
    
    # Sort inventory in descending order for greedy approach
    sorted_inv = sorted(inventory, reverse=True)
    
    while balls_sold < orders:
        # Find maximum value
        max_val = max(sorted_inv)
        max_idx = sorted_inv.index(max_val)
        
        print(f"Sell ball at price {max_val}")
        total_profit += max_val
        sorted_inv[max_idx] -= 1
        balls_sold += 1
    
    print(f"Total profit: {total_profit}")

detailed_example()

Selling process:
Sell ball at price 7
Sell ball at price 6
Sell ball at price 5
Sell ball at price 5
Sell ball at price 4
Sell ball at price 4
Total profit: 31

Conclusion

This problem uses binary search to efficiently find the optimal price threshold, combined with arithmetic series formulas to calculate maximum profit. The key insight is that we should always sell the highest-valued balls first to maximize total profit.