
思路无非每次找最大的即可
但是模拟超时
因此二分找到去到全都是num时候需要的球数
这个球数不能超过orders
否则就要更多球
然后多出来的那部分球就算是num的贡献
其余的等差数列求和即可
class Solution:
def maxProfit(self, inventory: List[int], orders: int) -> int:
MOD = 10 ** 9 + 7
inventory.sort()
n = len(inventory)
def get_balls(num):
# 所有大于num都变成num所需要的球数
idx = bisect_right(inventory, num) # 第一个大于num的
cnt = 0
for i in range(idx, n):
cnt += inventory[i] - num
return cnt
l, r = 0, inventory[-1]
while l + 1 < r:
mid = (l + r) // 2
if get_balls(mid) <= orders:
r = mid
else:
l = mid + 1
ball1 = get_balls(r)
ball2 = get_balls(l)
ballFinal, numFinal = -1, -1
if ball2 < orders:
ballFinal, numFinal = ball2, l
else:
ballFinal, numFinal = ball1, r
rest = orders - ballFinal
ans = 0
ans = numFinal * rest # rest
idx = bisect_right(inventory, numFinal)
for i in range(idx, n):
ans += ((numFinal + 1 + inventory[i]) * (inventory[i] - numFinal) // 2)
return ans % MOD
二分找边界 + 等差数列