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

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>

const int M = 1'000'000'007;

class Solution {
public:
    int maxProfit(std::vector<int>& inventory, int orders) {
        std::sort(inventory.begin(), inventory.end(), std::greater<int>());
        long long n = inventory.size();
        inventory.push_back(0);
        
        long long sum = 0;
        for (int i = 0; i < n; ++i) {
            long long diff = inventory[i] - inventory[i + 1];
            long long add = cal(inventory[i], inventory[i + 1]);
            long long num = diff * (i + 1);
            if (num < orders) {
                orders -= num;
                sum += add * (i + 1);
            }
            else {
                long long new_diff = orders / (i + 1);
                long long new_num = orders % (i + 1);
                sum += cal(inventory[i], inventory[i] - new_diff) * (i + 1);
                sum += new_num * (inventory[i] - new_diff);
                break;
            }
        }
        
        return sum % M;
    }
    
private:
    long long cal(long long a, long long b) {
        return (a + b + 1) * (a - b) / 2;
    }
};