Largest Values From Labels LeetCode solution in Java

1
class Solution {
2
  public int largestValsFromLabels(int[] values, int[] labels, int numWanted, int useLimit) {
3
    PriorityQueue<Pair<Integer, Integer>> maxHeap =
4
        new PriorityQueue<>((a, b) -> Integer.compare(b.getKey(), a.getKey()));
5
    for (int i = 0; i < values.length; i++) {
6
      maxHeap.add(new Pair<Integer, Integer>(values[i], labels[i]));
7
    }
8
    Map<Integer, Integer> labelLimitMap = new HashMap<>();
9
    int sum = 0;
10
    while (numWanted != 0) {
11
      int label = maxHeap.peek().getValue();
12
      if (labelLimitMap.containsKey(label)) {
13
        if (labelLimitMap.get(label) == useLimit) {
14
          maxHeap.poll();
15
        } else {
16
          labelLimitMap.put(label, labelLimitMap.get(label) + 1);
17
          sum += maxHeap.poll().getKey();
18
          numWanted--;
19
        }
20
      } else {
21
        labelLimitMap.put(label, 1);
22
        sum += maxHeap.poll().getKey();
23
        numWanted--;
24
      }
25
      // This Holds since at most numWanted is mentioned.
26
      if (maxHeap.isEmpty()) {
27
        return sum;
28
      }
29
    }
30
    return sum;
31
  }
32
}