Update Top K Frequent Elements.java

Author: varunu28

Medium/Top K Frequent Elements.java

@@ -4,21 +4,22 @@ public int[] topKFrequent(int[] nums, int k) {
     for (int num : nums) {
       map.put(num, map.getOrDefault(num, 0) + 1);
     }
-    PriorityQueue<Integer> pq = new PriorityQueue<>(new Comparator<Integer>(){
-      public int compare(Integer p1, Integer p2) {
-        return map.get(p1) - map.get(p2);
-      }
-    });
+    Map<Integer, List<Integer>> frequencyToValueMap = new HashMap<>();
     for (Integer key : map.keySet()) {
-      pq.add(key);
-      if (pq.size() > k) {
-        pq.poll();
-      }
+      frequencyToValueMap.computeIfAbsent(map.get(key), j -> new ArrayList<>()).add(key);
     }
-    int[] ans = new int[k];
-    for (int i = 0; i < k; i++) {
-      ans[i] = pq.poll();
+    List<Integer> result = new ArrayList<>();
+    for (int i = nums.length; i >= 0 && result.size() < k; i--) {
+      List<Integer> values = frequencyToValueMap.getOrDefault(i, new ArrayList<>());
+      if (result.size() + values.size() <= k) {
+        result.addAll(values);
+      } else {
+        int idx = 0;
+        while (result.size() < k) {
+          result.add(values.get(idx++));
+        }
+      }
     }
-    return ans;
+    return result.stream().mapToInt(i -> i).toArray();
   }
 }