Maximum Elegance of a K-Length Subsequence
Problem
You are given a 0-indexed 2D integer array items of length n and an integer k.
items[i] = [profiti, categoryi], where profiti and categoryi denote the profit and category of the ith item respectively.
Let's define the elegance of a subsequence of items as `total_profit
- distinct_categories2
, wheretotal_profitis the sum of all profits in the subsequence, anddistinct_categories` is the number of distinct categories from all the categories in the selected subsequence.
Your task is to find the maximum elegance from all subsequences of size
k in items.
Return an integer denoting the maximum elegance of a subsequence ofitems
with size exactlyk.
Note: A subsequence of an array is a new array generated from the original array by deleting some elements (possibly none) without changing the remaining elements' relative order.
Examples
Example 1
Input: items = [[3,2],[5,1],[10,1]], k = 2
Output: 17
Explanation: In this example, we have to select a subsequence of size 2.
We can select items[0] = [3,2] and items[2] = [10,1].
The total profit in this subsequence is 3 + 10 = 13, and the subsequence contains 2 distinct categories [2,1].
Hence, the elegance is 13 + 22 = 17, and we can show that it is the maximum achievable elegance.
Example 2
Input: items = [[3,1],[3,1],[2,2],[5,3]], k = 3
Output: 19
Explanation: In this example, we have to select a subsequence of size 3.
We can select items[0] = [3,1], items[2] = [2,2], and items[3] = [5,3].
The total profit in this subsequence is 3 + 2 + 5 = 10, and the subsequence contains 3 distinct categories [1,2,3].
Hence, the elegance is 10 + 32 = 19, and we can show that it is the maximum achievable elegance.
Example 3
Input: items = [[1,1],[2,1],[3,1]], k = 3
Output: 7
Explanation: In this example, we have to select a subsequence of size 3.
We should select all the items.
The total profit will be 1 + 2 + 3 = 6, and the subsequence contains 1 distinct category [1].
Hence, the maximum elegance is 6 + 12 = 7.
Constraints
1 <= items.length == n <= 10^5items[i].length == 2items[i][0] == profitiitems[i][1] == categoryi1 <= profiti <= 10^91 <= categoryi <= n1 <= k <= n
Solution
Method 1 – Greedy with Heap and HashSet
Intuition
To maximize elegance, we want to maximize both the total profit and the number of distinct categories. We start by picking the top k items by profit, then try to swap out duplicate categories for new ones (with lower profit) to increase the number of distinct categories, using a min-heap to track the smallest profit among duplicates.
Approach
- Sort items by profit in descending order.
- Select the first
kitems as the initial subsequence. - Track the total profit and the set of categories in the initial selection.
- Use a min-heap to keep profits of items with duplicate categories.
- For the rest of the items, if their category is new and the heap is not empty, swap the smallest duplicate profit for this new category to increase distinct categories.
- For each step, update the answer as
total_profit + (distinct_categories) ** 2. - Return the maximum elegance found.
Code
C++
class Solution {
public:
long long findMaximumElegance(vector<vector<int>>& items, int k) {
sort(items.rbegin(), items.rend());
long long ans = 0, sum = 0;
unordered_set<int> cats;
priority_queue<int, vector<int>, greater<int>> dup;
for (int i = 0; i < k; ++i) {
sum += items[i][0];
if (!cats.insert(items[i][1]).second) dup.push(items[i][0]);
}
ans = sum + 1LL * cats.size() * cats.size();
for (int i = k; i < items.size(); ++i) {
if (cats.count(items[i][1]) == 0 && !dup.empty()) {
sum += items[i][0] - dup.top();
dup.pop();
cats.insert(items[i][1]);
ans = max(ans, sum + 1LL * cats.size() * cats.size());
}
}
return ans;
}
};
Go
func findMaximumElegance(items [][]int, k int) int64 {
sort.Slice(items, func(i, j int) bool { return items[i][0] > items[j][0] })
sum := 0
cats := map[int]bool{}
dup := &IntHeap{}
heap.Init(dup)
for i := 0; i < k; i++ {
sum += items[i][0]
if cats[items[i][1]] {
heap.Push(dup, items[i][0])
}
cats[items[i][1]] = true
}
ans := int64(sum + len(cats)*len(cats))
for i := k; i < len(items); i++ {
if !cats[items[i][1]] && dup.Len() > 0 {
sum += items[i][0] - heap.Pop(dup).(int)
cats[items[i][1]] = true
ans = max(ans, int64(sum+len(cats)*len(cats)))
}
}
return ans
}
// IntHeap and max helper omitted for brevity
Java
class Solution {
public long findMaximumElegance(int[][] items, int k) {
Arrays.sort(items, (a, b) -> b[0] - a[0]);
long sum = 0, ans = 0;
Set<Integer> cats = new HashSet<>();
PriorityQueue<Integer> dup = new PriorityQueue<>();
for (int i = 0; i < k; i++) {
sum += items[i][0];
if (!cats.add(items[i][1])) dup.add(items[i][0]);
}
ans = sum + 1L * cats.size() * cats.size();
for (int i = k; i < items.length; i++) {
if (!cats.contains(items[i][1]) && !dup.isEmpty()) {
sum += items[i][0] - dup.poll();
cats.add(items[i][1]);
ans = Math.max(ans, sum + 1L * cats.size() * cats.size());
}
}
return ans;
}
}
Kotlin
class Solution {
fun findMaximumElegance(items: Array<IntArray>, k: Int): Long {
items.sortByDescending { it[0] }
var sum = 0L
val cats = mutableSetOf<Int>()
val dup = java.util.PriorityQueue<Int>()
for (i in 0 until k) {
sum += items[i][0]
if (!cats.add(items[i][1])) dup.add(items[i][0])
}
var ans = sum + cats.size.toLong() * cats.size
for (i in k until items.size) {
if (items[i][1] !in cats && dup.isNotEmpty()) {
sum += items[i][0] - dup.poll()
cats.add(items[i][1])
ans = maxOf(ans, sum + cats.size.toLong() * cats.size)
}
}
return ans
}
}
Python
import heapq
class Solution:
def findMaximumElegance(self, items: list[list[int]], k: int) -> int:
items.sort(reverse=True)
s = 0
cats = set()
dup = []
for i in range(k):
s += items[i][0]
if items[i][1] in cats:
heapq.heappush(dup, items[i][0])
cats.add(items[i][1])
ans = s + len(cats) * len(cats)
for i in range(k, len(items)):
if items[i][1] not in cats and dup:
s += items[i][0] - heapq.heappop(dup)
cats.add(items[i][1])
ans = max(ans, s + len(cats) * len(cats))
return ans
Rust
impl Solution {
pub fn find_maximum_elegance(items: Vec<Vec<i32>>, k: i32) -> i64 {
let k = k as usize;
let mut items = items;
items.sort_by(|a, b| b[0].cmp(&a[0]));
let mut sum = 0i64;
let mut cats = std::collections::HashSet::new();
let mut dup = std::collections::BinaryHeap::new();
for i in 0..k {
sum += items[i][0] as i64;
if !cats.insert(items[i][1]) {
dup.push(std::cmp::Reverse(items[i][0]));
}
}
let mut ans = sum + (cats.len() * cats.len()) as i64;
for i in k..items.len() {
if !cats.contains(&items[i][1]) && !dup.is_empty() {
sum += items[i][0] as i64 - dup.pop().unwrap().0 as i64;
cats.insert(items[i][1]);
ans = ans.max(sum + (cats.len() * cats.len()) as i64);
}
}
ans
}
}
TypeScript
class Solution {
findMaximumElegance(items: number[][], k: number): number {
items.sort((a, b) => b[0] - a[0]);
let sum = 0, ans = 0;
const cats = new Set<number>();
const dup: number[] = [];
for (let i = 0; i < k; i++) {
sum += items[i][0];
if (cats.has(items[i][1])) dup.push(items[i][0]);
cats.add(items[i][1]);
}
dup.sort((a, b) => a - b);
ans = sum + cats.size * cats.size;
let idx = 0;
for (let i = k; i < items.length; i++) {
if (!cats.has(items[i][1]) && dup.length > idx) {
sum += items[i][0] - dup[idx++];
cats.add(items[i][1]);
ans = Math.max(ans, sum + cats.size * cats.size);
}
}
return ans;
}
}
Complexity
- ⏰ Time complexity:
O(n log n), wherenis the number of items, due to sorting and heap operations. - 🧺 Space complexity:
O(n), for storing sets and heaps.