Maximize Sum of At Most K Distinct Elements

Difficulty: Easy | Acceptance: 77.00% | Paid: No Topics: Array, Hash Table, Greedy, Sorting

You are given an integer array nums and an integer k. You need to choose a subset of elements from nums such that the number of distinct elements in the subset is at most k. Return the maximum possible sum of the elements in the chosen subset.

• Examples
• Constraints
• Greedy with Sorting
• Greedy with Max Heap

Examples

Example 1:

Input: nums = [1,2,3,3,4,4,4], k = 2
Output: 18
Explanation: We can choose elements 4 and 3. There are three 4s and two 3s. Sum = 4*3 + 3*2 = 18.

Example 2:

Input: nums = [5,5,5], k = 1
Output: 15
Explanation: We choose all three 5s. The number of distinct elements is 1, which is <= k.

Example 3:

Input: nums = [1,2,3], k = 0
Output: 0
Explanation: We cannot choose any distinct elements, so the sum is 0.

Constraints

- 1 <= nums.length <= 100
- 1 <= nums[i] <= 10^9
- 1 <= k <= nums.length

Greedy with Sorting

Intuition To maximize the sum, we should prioritize the largest numbers available in the array. Since we are allowed to take all occurrences of a chosen distinct number, the optimal strategy is to select the k largest distinct values and sum all their occurrences.

Steps

• Count the frequency of each number using a hash map.
• Extract the distinct numbers and sort them in descending order.
• Iterate through the sorted list, summing the value multiplied by its frequency.
• Stop after processing k distinct numbers.

from collections import Counter
from typing import List

class Solution:
    def maxSum(self, nums: List[int], k: int) -> int:
        if k == 0:
            return 0
        
        freq = Counter(nums)
        # Sort distinct numbers in descending order
        distinct_nums = sorted(freq.keys(), reverse=True)
        
        total = 0
        # Take top k distinct numbers
        for i in range(min(k, len(distinct_nums))):
            num = distinct_nums[i]
            total += num * freq[num]
            
        return total

Complexity

• Time: O(N log N) due to sorting the distinct elements.
• Space: O(N) to store the frequency map.
• Notes: Sorting is the most straightforward approach and is efficient enough given the constraints.

Greedy with Max Heap

Intuition Instead of sorting all distinct elements, we can use a Max Heap (Priority Queue) to repeatedly extract the largest remaining distinct number. This is useful if we want to avoid sorting the entire list when k is much smaller than the number of distinct elements.

Steps

• Count the frequency of each number.
• Insert all distinct numbers into a Max Heap.
• Repeat k times: extract the maximum value from the heap, add value * frequency to the sum, and decrement k.

import heapq
from collections import Counter
from typing import List

class Solution:
    def maxSum(self, nums: List[int], k: int) -> int:
        if k == 0:
            return 0
        
        freq = Counter(nums)
        # Python heapq is a min-heap, so we invert values for max-heap behavior
        max_heap = [-num for num in freq.keys()]
        heapq.heapify(max_heap)
        
        total = 0
        while k > 0 and max_heap:
            num = -heapq.heappop(max_heap)
            total += num * freq[num]
            k -= 1
            
        return total

Complexity

• Time: O(N + D log D) to build the heap, where D is the number of distinct elements. Extracting k elements takes O(k log D).
• Space: O(D) for the heap and frequency map.
• Notes: This approach is more efficient than sorting if k is significantly smaller than D, though in the worst case (k close to D), it is comparable to sorting.