347. Top K Frequent Elements

medium arrays-hashing

Return the k most frequent elements.

Given an integer array nums and an integer k, return the k most frequent elements.

Approach

Count frequencies, then use a bucket sort approach (index = frequency) — O(n) time, O(n) space.

Solutions

Ruby

def top_k_frequent(nums, k)
  counts = nums.tally
  bucket = Array.new(nums.size + 1) { [] }
  counts.each { |num, freq| bucket[freq] << num }
  bucket.reverse.flatten.first(k)
end

Python

from collections import Counter

def top_k_frequent(nums: list[int], k: int) -> list[int]:
    counts = Counter(nums)
    bucket = [[] for _ in range(len(nums) + 1)]
    for num, freq in counts.items():
        bucket[freq].append(num)
    result = []
    for i in range(len(bucket) - 1, -1, -1):
        result.extend(bucket[i])
        if len(result) >= k:
            break
    return result[:k]

Go

func topKFrequent(nums []int, k int) []int {
    counts := make(map[int]int)
    for _, n := range nums {
        counts[n]++
    }
    bucket := make([][]int, len(nums)+1)
    for num, freq := range counts {
        bucket[freq] = append(bucket[freq], num)
    }
    result := []int{}
    for i := len(bucket) - 1; i >= 0 && len(result) < k; i-- {
        result = append(result, bucket[i]...)
    }
    return result[:k]
}

TypeScript

function topKFrequent(nums: number[], k: number): number[] {
    const counts = new Map<number, number>();
    for (const n of nums) counts.set(n, (counts.get(n) ?? 0) + 1);
    const bucket: number[][] = Array.from({ length: nums.length + 1 }, () => []);
    for (const [num, freq] of counts) bucket[freq].push(num);
    const result: number[] = [];
    for (let i = bucket.length - 1; i >= 0 && result.length < k; i--) {
        result.push(...bucket[i]);
    }
    return result.slice(0, k);
}

Elixir

defmodule Solution do
  def top_k_frequent(nums, k) do
    nums
    |> Enum.frequencies()
    |> Enum.group_by(fn {_num, freq} -> freq end, fn {num, _freq} -> num end)
    |> Enum.sort_by(fn {freq, _nums} -> freq end, :desc)
    |> Enum.flat_map(fn {_freq, nums} -> nums end)
    |> Enum.take(k)
  end
end