Maximum Total Beauty of the Gardens LeetCode solution in C++

1
class Solution {
2
 public:
3
  long long maximumBeauty(vector<int> &flowers, long long newFlowers, int target, int full,
4
                          int partial) {
5
    long long ans = 0;
6
    int n = flowers.size();
7
    if (n == 0) return 0;
8
    sort(flowers.begin(), flowers.end());
9
    vector<long long> pref(n, 0); // Prefix sums array
10
    pref[0] = flowers[0];
11
    for (int i = 1; i < n; i++) {
12
      pref[i] = pref[i - 1] + flowers[i];
13
    }
14
    for (int i = n; i >= 0; i--) {
15
      // When i is n, none of the elements are incremented till target.
16
      if (i < n) { // The ith element(0-indexed) is incremented till target. This is
17
                   // handling the case of "full". We are doing this from the end
18
                   // because it will cost the least number of new Flowers.
19
        long long toAdd = max(0, target - flowers[i]);
20
        if (newFlowers < toAdd) {
21
          break; // We are breaking because we would've already found the best
22
                 // answer in the previous iterations
23
        }
24
        newFlowers -= toAdd;
25
        flowers.pop_back();
26
      }
27
      long long beautyForFullCase =
28
          (long long)(n - i) * full; // All the elements from i to n -1 are incremented till target,
29
                                     // so they have a beauty of (n - 1) * full
30

31
      // Here, we do binary search on x, which is the the maximum number we can
32
      // increment all the numbers from 0 to i - 1 to. The maximum we can
33
      // increment to is target-1 because after that, "partial" will not be
34
      // applicable. Binary search is done to find out the maximum possible
35
      // answer for the partial case.
36
      int l = 0, r = target - 1;
37
      long long beautyForPartialCase = 0;
38
      while (l <= r) {
39
        long long mid = l + (r - l) / 2; // mid is the maximum number we can increment to
40
        auto it = upper_bound(flowers.begin(), flowers.end(),
41
                              mid);  // will point to an element that is atleast 1 more than the
42
                                     // incomplete garden with most number of flowers
43
        if (it == flowers.begin()) { // there are no incomplete gardens
44
          l = mid + 1;
45
        } else {
46
          long long len = distance(flowers.begin(), it); // number of incomplete gardens
47
          long long sum = pref[len - 1];
48
          long long needed = ((long long)mid * len) - sum; // number of flowers needed to make
49
                                                           // "mid" as the minimum element
50
          if (needed <= newFlowers) { // if it is possible to plant flowers such that the
51
                                      // minimum number of flowers for any garden is
52
                                      // "mid", then we go to check for higher values
53
                                      // through Binary search
54
            l = mid + 1;
55
            beautyForPartialCase =
56
                (long long)mid * partial; // (minimum number of flowers in any incomplete garden)
57
                                          // * (partial) = total beauty of incomplete gardens
58
          } else {
59
            r = mid - 1;
60
          }
61
        }
62
      }
63
      ans =
64
          max(ans, beautyForFullCase + beautyForPartialCase); // checking if this answer is better
65
                                                              // than the previous answers computed.
66
    }
67
    return ans;
68
  }
69
};