Compute the k closest stars

Star drawing

Compute the k closest stars to earth (position 0,0,0) that are stored in a large file. The file is too large to store in memory.

The normal brute force way is to load everything and sort the to the k closest. But since we can’t do that we need a way to store just the k number of elements and evict large values

max heap

Putting this together in code:

import heapq

class Star:
    def __init__(self, x, y, z):
        self.x, self.y self.z = x, y, z

def closest_k_stars(stars, k):
    heap = []

    for star in stars:
        # no max heap supported in python so we can just use negative values to mimic the max heap behavior
        heapq.heappush(heap, (-star.distance, star))

        if len(heap) > k:
            heapq.heappop(heap)

    return [x[1] for x in heapq.nlargest(k, heap)]