2813. Maximum Elegance of a K-Length Subsequence

link

In contest, I alrady think of maximum profit and heap.

However, I didn't solve it.

class Solution:
    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:
        
        n = len(items)
        
        items.sort(key=lambda x: -x[0])
        
        M = defaultdict(list)
        
        s = set()

        h = []

        res = 0
        total = 0
        for i in range(n):
            if i < k:
                total += items[i][0]

                if items[i][1] not in s:
                    s.add(items[i][1])
                else:
                    heappush(h, items[i])
            else:
                if h and items[i][1] not in s:
                    nxt = heappop(h)
                    total -= nxt[0]
                    total += items[i][0]
                    s.add(items[i][1])

            res = max(res, total + len(s) ** 2)
        
        return res