Top ‘K’ Elements (Heaps)

1. Kth Largest Element in an Array Medium

Problem

Return the k-th largest element in an unsorted array. Note this is the k-th largest in sorted order, not the k-th distinct value.

Approach

• Keep a min-heap of exactly the k largest values seen. When the scan ends, the root is the smallest of the top k, which is precisely the k-th largest.
• Seed the heap with the first k values and heapify them in O(k), rather than pushing them one at a time.
• Be ready with quickselect as the O(n)-average alternative, and with sorted(nums)[-k] as the one-liner.

Solution: heap

import heapq


def find_kth_largest(nums: list[int], k: int) -> int:
    """The k-th largest value in nums, counting duplicates.

    Args:
        nums: Integers, unsorted.
        k: 1-based rank from the top, with 1 <= k <= len(nums).

    Returns:
        The k-th largest value.

    Raises:
        ValueError: If k is out of range.

    Example:
        >>> find_kth_largest([3, 2, 1, 5, 6, 4], 2)
        5
    """
    if not 1 <= k <= len(nums):
        raise ValueError(f"k must be in 1..{len(nums)}, got {k}")

    # Seed with the first k values. heapify is O(k), k pushes would be O(k log k).
    heap = nums[:k]
    heapq.heapify(heap)

    for value in nums[k:]:
        if value > heap[0]:
            heapq.heapreplace(heap, value)

    # The root is the weakest of the k largest, which is the k-th largest.
    return heap[0]

Solution: quickselect

import random


def _partition(nums: list[int], left: int, right: int, pivot_index: int) -> int:
    """Lomuto partition. Returns the pivot's final resting index."""
    pivot = nums[pivot_index]
    nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

    store = left
    for i in range(left, right):
        if nums[i] < pivot:
            nums[store], nums[i] = nums[i], nums[store]
            store += 1

    nums[store], nums[right] = nums[right], nums[store]
    return store


def find_kth_largest_quickselect(nums: list[int], k: int) -> int:
    """O(n) expected time. Mutates nums.

    The k-th largest sits at index len(nums) - k once the array is
    partitioned around it, so we only recurse into the half that contains
    that index.
    """
    target = len(nums) - k
    left, right = 0, len(nums) - 1

    while left <= right:
        # A random pivot is what turns the O(n^2) worst case into a
        # vanishingly unlikely one.
        pivot_index = _partition(nums, left, right, random.randint(left, right))

        if pivot_index == target:
            return nums[pivot_index]
        if pivot_index < target:
            left = pivot_index + 1
        else:
            right = pivot_index - 1

    raise ValueError("k out of range")

Which to offer

Edge cases to raise

• k == 1: the maximum. k == len(nums): the minimum.
• Duplicates: [3, 3, 3] with k = 2 gives 3. Rank counts positions, not distinct values.
• Negative values: nothing changes.
• Quickselect mutates the input. Say so, or copy first.

2. Top K Frequent Elements Medium

Problem

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

Approach

• Two phases. Count, then select. The counting is a Counter in O(n). The selection is where the choice lives.
• Heap: keep the k highest counts, O(m log k) where m is the number of distinct values.
• Bucket sort: a count can never exceed n, so index an array of lists by count and walk it downward. That is O(n) and beats the heap. Interviewers who ask for “better than O(n log k)” are asking for this.

Solution: heap

import heapq
from collections import Counter


def top_k_frequent(nums: list[int], k: int) -> list[int]:
    """The k most frequent values in nums, in no particular order.

    Args:
        nums: Input values.
        k: How many to return, with 1 <= k <= number of distinct values.

    Returns:
        The k values with the highest counts.

    Example:
        >>> sorted(top_k_frequent([1, 1, 1, 2, 2, 3], 2))
        [1, 2]
    """
    counts = Counter(nums)
    heap: list[tuple[int, int]] = []          # (frequency, value)

    for value, frequency in counts.items():
        if len(heap) < k:
            heapq.heappush(heap, (frequency, value))
        elif frequency > heap[0][0]:
            heapq.heapreplace(heap, (frequency, value))

    return [value for _, value in heap]

Solution: bucket by count, O(n)

from collections import Counter


def top_k_frequent_buckets(nums: list[int], k: int) -> list[int]:
    """Same answer in linear time, by indexing buckets on the count itself.

    A value cannot appear more than len(nums) times, so counts are bounded
    and can be used directly as array indices. That removes the log factor.
    """
    counts = Counter(nums)

    # buckets[f] holds every value that occurs exactly f times.
    buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]
    for value, frequency in counts.items():
        buckets[frequency].append(value)

    out: list[int] = []
    for frequency in range(len(nums), 0, -1):
        for value in buckets[frequency]:
            out.append(value)
            if len(out) == k:
                return out

    return out

Walkthrough of the bucket version

nums = [1, 1, 1, 2, 2, 3], k = 2. Counts are {1: 3, 2: 2, 3: 1}. Buckets: index 1 holds [3], index 2 holds [2], index 3 holds [1]. Walking down from index 6: nothing at 6, 5, 4; at 3 take 1; at 2 take 2; we now have two, so return [1, 2].

The production one-liner

def top_k_frequent_stdlib(nums: list[int], k: int) -> list[int]:
    """What you would actually ship."""
    return [value for value, _ in Counter(nums).most_common(k)]

most_common uses heapq.nlargest internally when k is given, so it is the heap solution with none of the code. Write your own first, then show this.

Edge cases to raise

• k equals the number of distinct values: returns all of them.
• Ties in frequency: the problem says any order, so either value may win. Confirm that.
• Every value unique: all counts are 1, and the first k encountered are returned.
• Empty input with k = 0: returns []. Ask whether k can be zero.

3. K Closest Points to Origin Medium

Problem

Given a list of points on a plane and an integer k, return the k points closest to the origin. The distance is the usual Euclidean distance.

Two observations before any code

• Do not compute square roots. sqrt is monotone increasing, so ordering by x² + y² gives exactly the same ranking as ordering by √(x² + y²). Skipping it removes n floating-point operations and, more importantly, keeps the comparisons in exact integer arithmetic. Say this unprompted.
• We want the K smallest distances, so the size-K heap must be a max-heap. Python only has a min-heap, so negate the distance. The root is then the furthest point currently kept, which is again the eviction candidate.

Solution

import heapq


def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    """The k points nearest the origin, in no particular order.

    Args:
        points: [x, y] integer coordinates.
        k: How many to return, with 1 <= k <= len(points).

    Returns:
        The k closest points.

    Example:
        >>> sorted(k_closest([[1, 3], [-2, 2], [5, 8], [0, 1]], 2))
        [[-2, 2], [0, 1]]
    """
    # Negated squared distance turns heapq's min-heap into a max-heap, so the
    # root is the furthest point we are currently keeping. x and y are the
    # tie-break, which keeps every comparison on plain integers.
    heap: list[tuple[int, int, int]] = []

    for x, y in points:
        entry = (-(x * x + y * y), x, y)

        if len(heap) < k:
            heapq.heappush(heap, entry)
        elif entry[0] > heap[0][0]:      # closer than the worst one kept
            heapq.heapreplace(heap, entry)

    return [[x, y] for _, x, y in heap]

Walkthrough

points = [[1,3],[-2,2],[5,8],[0,1]], k = 2. Squared distances are 10, 8, 89, 1.

Result: [[-2, 2], [0, 1]].

The alternatives, ranked

Edge cases to raise

• k == len(points): every point is returned.
• Duplicate points, or ties in distance: any tied point may be chosen, and the (d, x, y) tuple makes the comparison total so nothing raises.
• Points at the origin: distance 0, always closest.
• Large coordinates: squaring can overflow in a fixed-width language. Python integers cannot, but mention it.