Skip to content

1648. Sell Diminishing-Valued Colored Balls

  • Time: $O(\texttt{sort})$
  • Space: $O(\texttt{sort})$
 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
class Solution {
 public:
  int maxProfit(vector<int>& inventory, int orders) {
    constexpr int kMod = 1'000'000'007;
    long ans = 0;
    long largestCount = 1;

    ranges::sort(inventory, greater<>());

    for (int i = 0; i < inventory.size(); ++i, ++largestCount)
      if (i == inventory.size() - 1 || inventory[i] > inventory[i + 1]) {
        // If we are at the last inventory, or inventory[i] > inventory[i + 1].
        // In either case, we will pick inventory[i - largestCount + 1..i].
        const int pick = (i == inventory.size() - 1)
                             ? inventory[i]
                             : inventory[i] - inventory[i + 1];
        if (largestCount * pick >= orders) {
          // We have run out of orders, so we need to recalculate the number of
          // balls that we actually pick for inventory[i - largestCount + 1..i].
          const int actualPick = orders / largestCount;
          const int remaining = orders % largestCount;
          return (ans +
                  largestCount *
                      trapezoid(inventory[i], inventory[i] - actualPick + 1) +
                  remaining * static_cast<long>(inventory[i] - actualPick)) %
                 kMod;
        }
        ans += largestCount * trapezoid(inventory[i], inventory[i] - pick + 1);
        ans %= kMod;
        orders -= largestCount * pick;
      }

    throw;
  }

 private:
  long trapezoid(long a, long b) {
    return (a + b) * (a - b + 1) / 2;
  }
};

 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
class Solution {
  public int maxProfit(int[] inventory, int orders) {
    final int kMod = 1_000_000_007;
    long ans = 0;
    long largestCount = 1;

    Arrays.sort(inventory);

    for (int i = inventory.length - 1; i >= 0; --i, ++largestCount)
      if (i == 0 || inventory[i] > inventory[i - 1]) {
        // If we are at the first inventory, or inventory[i] > inventory[i + 1].
        // In either case, we will pick inventory[i..i + largestCount - 1].
        final long pick = (i == 0) ? inventory[i] : inventory[i] - inventory[i - 1];
        if (largestCount * pick >= orders) {
          // We have run out of orders, so we need to recalculate the number of
          // balls that we actually pick for inventory[i..i + largestCount - 1].
          final long actualPick = orders / largestCount;
          final long remaining = orders % largestCount;
          return (int) ((ans +
                         largestCount * trapezoid(inventory[i], inventory[i] - actualPick + 1) +
                         remaining * (inventory[i] - actualPick)) %
                        kMod);
        }
        ans += largestCount * trapezoid(inventory[i], inventory[i] - pick + 1);
        ans %= kMod;
        orders -= largestCount * pick;
      }

    throw new IllegalArgumentException();
  }

  private long trapezoid(long a, long b) {
    return (a + b) * (a - b + 1) / 2;
  }
}

 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
class Solution:
  def maxProfit(self, inventory: List[int], orders: int) -> int:
    kMod = 1_000_000_007
    ans = 0
    largestCount = 1

    def trapezoid(a: int, b: int) -> int:
      return (a + b) * (a - b + 1) // 2

    for a, b in itertools.pairwise(sorted(inventory, reverse=True) + [0]):
      if a > b:
        # If we are at the last inventory, or inventory[i] > inventory[i + 1].
        # In either case, we will pick inventory[i - largestCount + 1..i].
        pick = a - b
        # We have run out of orders, so we need to recalculate the number of
        # balls that we actually pick for inventory[i - largestCount + 1..i].
        if largestCount * pick >= orders:
          actualPick, remaining = divmod(orders, largestCount)
          return (ans +
                  largestCount * trapezoid(a, a - actualPick + 1) +
                  remaining * (a - actualPick)) % kMod
        ans += largestCount * trapezoid(a, a - pick + 1)
        ans %= kMod
        orders -= largestCount * pick
      largestCount += 1