Skip to content

2813. Maximum Elegance of a K-Length Subsequence 👍

  • Time: $O(\texttt{sort})$
  • Space: $O(n)$
 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
41
42
43
44
45
46
class Solution {
 public:
  long long findMaximumElegance(vector<vector<int>>& items, int k) {
    long ans = 0;
    long totalProfit = 0;
    unordered_set<int> seenCategories;
    stack<int> decreasingDuplicateProfits;

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

    for (int i = 0; i < k; i++) {
      const int profit = items[i][0];
      const int category = items[i][1];
      totalProfit += profit;
      if (seenCategories.contains(category))
        decreasingDuplicateProfits.push(profit);
      else
        seenCategories.insert(category);
    }

    ans = totalProfit +
          static_cast<long>(seenCategories.size()) * seenCategories.size();

    for (int i = k; i < items.size(); ++i) {
      const int profit = items[i][0];
      const int category = items[i][1];
      if (!seenCategories.contains(category) &&
          !decreasingDuplicateProfits.empty()) {
        // If this is a new category we haven't seen before, it's worth
        // considering taking it and replacing the one with the least profit
        // since it will increase the distinct_categories and potentially result
        // in a larger total_profit + distinct_categories^2.
        totalProfit -= decreasingDuplicateProfits.top(),
            decreasingDuplicateProfits.pop();
        totalProfit += profit;
        seenCategories.insert(category);
        ans = max(ans,
                  static_cast<long>(totalProfit +
                                    static_cast<long>(seenCategories.size()) *
                                        seenCategories.size()));
      }
    }

    return ans;
  }
};

 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
class Solution {
  public long findMaximumElegance(int[][] items, int k) {
    long ans = 0;
    long totalProfit = 0;
    Set<Integer> seenCategories = new HashSet<>();
    Deque<Integer> decreasingDuplicateProfits = new ArrayDeque<>();

    Arrays.sort(items, Comparator.comparingInt(item -> - item[0]));

    for (int i = 0; i < k; ++i) {
      final int profit = items[i][0];
      final int category = items[i][1];
      totalProfit += profit;
      if (seenCategories.contains(category))
        decreasingDuplicateProfits.push(profit);
      else
        seenCategories.add(category);
    }

    ans = totalProfit + 1L * seenCategories.size() * seenCategories.size();

    for (int i = k; i < items.length; ++i) {
      final int profit = items[i][0];
      final int category = items[i][1];
      if (!seenCategories.contains(category) && !decreasingDuplicateProfits.isEmpty()) {
        // If this is a new category we haven't seen before, it's worth
        // considering taking it and replacing the one with the least profit
        // since it will increase the distinct_categories and potentially result
        // in a larger total_profit + distinct_categories^2.
        totalProfit -= decreasingDuplicateProfits.pop();
        totalProfit += profit;
        seenCategories.add(category);
        ans = Math.max(ans, totalProfit + 1L * seenCategories.size() * seenCategories.size());
      }
    }

    return ans;
  }
}

 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
class Solution:
  def findMaximumElegance(self, items: list[list[int]], k: int) -> int:
    ans = 0
    totalProfit = 0
    seenCategories = set()
    decreasingDuplicateProfits = []

    items.sort(reverse=True)

    for i in range(k):
      profit, category = items[i]
      totalProfit += profit
      if category in seenCategories:
        decreasingDuplicateProfits.append(profit)
      else:
        seenCategories.add(category)

    ans = totalProfit + len(seenCategories)**2

    for i in range(k, len(items)):
      profit, category = items[i]
      if category not in seenCategories and decreasingDuplicateProfits:
        # If this is a new category we haven't seen before, it's worth
        # considering taking it and replacing the one with the least profit
        # since it will increase the distinct_categories and potentially result
        # in a larger total_profit + distinct_categories^2.
        totalProfit -= decreasingDuplicateProfits.pop()
        totalProfit += profit
        seenCategories.add(category)
        ans = max(ans, totalProfit + len(seenCategories)**2)

    return ans