Maximum Elegance of a K-Length Subsequence

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

class Solution:
    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:
        items = sorted(items, key = lambda v: -v[0])
        result = 0
        current = 0
        stack = []
        seen = set()
        for i, (profit, category) in enumerate(items):
            if i < k:
                if category in seen:
                    stack.append(profit)
                current += profit
            elif category not in seen:
                if not stack:
                    break
                current += profit - stack.pop()
            seen.add(category)
            result = max(result, current + len(seen) * len(seen))
        return result