<p>You are given a <strong>0-indexed</strong> 2D integer array <code>items</code> of length <code>n</code> and an integer <code>k</code>.</p>

<p><code>items[i] = [profit<sub>i</sub>, category<sub>i</sub>]</code>, where <code>profit<sub>i</sub></code> and <code>category<sub>i</sub></code> denote the profit and category of the <code>i<sup>th</sup></code> item respectively.</p>

<p>Let&#39;s define the <strong>elegance</strong> of a <strong>subsequence</strong> of <code>items</code> as <code>total_profit + distinct_categories<sup>2</sup></code>, where <code>total_profit</code> is the sum of all profits in the subsequence, and <code>distinct_categories</code> is the number of <strong>distinct</strong> categories from all the categories in the selected subsequence.</p>

<p>Your task is to find the <strong>maximum elegance</strong> from all subsequences of size <code>k</code> in <code>items</code>.</p>

<p>Return <em>an integer denoting the maximum elegance of a subsequence of </em><code>items</code><em> with size exactly </em><code>k</code>.</p>

<p><strong>Note:</strong> 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&#39; relative order.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> items = [[3,2],[5,1],[10,1]], k = 2
<strong>Output:</strong> 17
<strong>Explanation: </strong>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 + 2<sup>2</sup> = 17, and we can show that it is the maximum achievable elegance. 
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> items = [[3,1],[3,1],[2,2],[5,3]], k = 3
<strong>Output:</strong> 19
<strong>Explanation:</strong> 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 + 3<sup>2</sup> = 19, and we can show that it is the maximum achievable elegance.</pre>

<p><strong class="example">Example 3:</strong></p>

<pre>
<strong>Input:</strong> items = [[1,1],[2,1],[3,1]], k = 3
<strong>Output:</strong> 7
<strong>Explanation:</strong> 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 + 1<sup>2</sup> = 7.  </pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= items.length == n &lt;= 10<sup>5</sup></code></li>
	<li><code>items[i].length == 2</code></li>
	<li><code>items[i][0] == profit<sub>i</sub></code></li>
	<li><code>items[i][1] == category<sub>i</sub></code></li>
	<li><code>1 &lt;= profit<sub>i</sub> &lt;= 10<sup>9</sup></code></li>
	<li><code>1 &lt;= category<sub>i</sub> &lt;= n </code></li>
	<li><code>1 &lt;= k &lt;= n</code></li>
</ul>

class Solution:
    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:
        # We can maintain priority queues for all categories each, and use another priority queue to maintain the element of the maximum profit for the respective categories.
        # If we do not take care about category factor, we can choose elements with top-k profits. If the current category number is K, we claim that the optimal solution must not have categories less than K.
        # To examine the optimal answer with K + 1 categories, we can remove the item with the smallest profit and add back the maximum one with the next category. Doing this until all the categories are traversed or there is only one item per category for the picked k elements.
        items.sort(key=lambda x:-x[0])
        fixed = set()
        flex = collections.defaultdict(list)
        tp = 0
        cat = 0
        ans = 0
        for i in range(k):
            if items[i][1] not in fixed:
                fixed.add(items[i][1])
                tp += items[i][0]
                cat += 1
            else:
                flex[items[i][1]].append(items[i][0])
                tp += items[i][0]
        ans = tp + cat * cat
        reducible = [(flex[k][-1], k) for k in flex]
        if len(reducible) == 0:
            return ans
        heapq.heapify(reducible)
        for i in range(k, len(items)):
            if items[i][1] not in fixed:
                # A new category is found
                fixed.add(items[i][1])
                cat += 1
                tp += items[i][0]
                # Prune the item with the smallest profit
                prof, cat_id = reducible[0]
                tp -= prof
                flex[cat_id].pop()
                if len(flex[cat_id]) > 0:
                    heapq.heapreplace(reducible, (flex[cat_id][-1], cat_id))
                else:
                    heapq.heappop(reducible)
                # Update the answer if applicable
                ans = max(ans, tp + cat * cat)
                if len(reducible) == 0:
                    return ans
        return ans
1280