Maximize total catch over K hours from ponds whose yield drops after each hour

There are N fishing ponds. M[i] is the number of fish you catch if you spend the next hour fishing at pond i. You have K hours in total. Each hour you pick any one pond and fish there for the whole hour. Every hour you spend at a pond lowers that pond's hourly catch by 1 for the following hours; ponds you do not fish at keep their current catch.

Return the maximum total number of fish you can catch in the K hours.

The interviewer required a solution whose running time is asymptotically better than $O(K \log N)$, so simulating the hours one at a time is too slow when K is large.

Function Signature

def max_fish(M: list[int], K: int) -> int:

Rules

• An hour at pond i catches its current hourly catch, then lowers that catch by 1.
• A pond's hourly catch never goes below 0: an hour at a pond whose catch is 0 catches nothing, and its catch stays 0.
• The order in which you visit ponds is free, and you may stay at the same pond for consecutive hours.
• Return only the maximum total, which is unique.

Constraints

• N == len(M) and 1 <= N <= 10^5
• 0 <= M[i] <= 10^6
• 0 <= K <= 10^9
• The answer is at most 10^15, which exceeds 2^31 - 1; Python integers handle this, but other languages need 64-bit integers.

Examples

Example 1

Input: M = [90, 100], K = 100
Output: 7075

One optimal plan spends 55 hours at the second pond (4015 fish) and 45 hours at the first pond (3060 fish).

Example 2

Input: M = [3, 5], K = 4
Output: 15

For example: two hours at the second pond (5 + 4), after which both ponds yield 3, then one hour at each (3 + 3).

Example 3

Input: M = [2], K = 5
Output: 3

The only pond yields 2, then 1, then 0 for each of the remaining three hours.