Top K Candidates with Bounded Max-Heap

The Problem with Naive Selection

In your current brute-force search, you iterate through all N vectors in the database, calculate N distances, and then need to identify the k nearest neighbors. The most straightforward way to do this is:

1. Store all N distances (along with their corresponding vector IDs) in an array.
2. Sort this entire array of N items based on distance.
3. Take the first k elements from the sorted array.

While simple, this approach has a significant performance bottleneck. The distance calculation is O(N), but the sort operation is typically O(N log N). For a large dataset, this sorting step will dominate the runtime and become prohibitively slow. We need a way to find the top k elements without sorting all N of them.

The Heap: An Efficient Solution

The ideal data structure for this task is a heap. A heap is a specialized tree-based data structure that satisfies the heap property: in a max-heap, for any given node, its value is greater than or equal to the values of its children. This means the largest element is always at the root. Conversely, in a min-heap, the root holds the smallest element.

Heaps are typically implemented using a simple array, which makes them very cache-friendly. The parent-child relationships are calculated using array indices, avoiding the need for pointers and the associated memory overhead and potential cache misses.

The key insight for our k-NN problem is to use a bounded max-heap of size k to find the k smallest distances. This may seem counter-intuitive, but the logic is highly effective.

The following video discusses this pattern. Pay close attention to the example of finding the k closest points to the origin, as it's a direct analogy to our problem of finding the k smallest distances.

The Bounded Max-Heap Algorithm

Here is the step-by-step algorithm for finding the k nearest neighbors using a max-heap of size k:

1. Initialize: Create a max-heap with a capacity of k.
2. Fill the heap: Iterate through the first k vectors in your database. For each one, calculate its distance to the query vector and push the result (distance and vector ID) onto the heap.
3. Iterate and Compare: For the remaining N-k vectors:
   a. Calculate the distance of the current vector to the query.
   b. Look at the root of the max-heap. The root holds the largest distance among the k nearest neighbors found so far.
   c. If the new distance is greater than or equal to the distance at the root, do nothing. This vector is not a better candidate than what we already have. This is a very fast O(1) rejection.
   d. If the new distance is smaller than the distance at theroot, it means we've found a new, closer neighbor. We must pop the root (removing the farthest neighbor from our top k set) and push the new, closer neighbor onto the heap. This combined pop-and-push operation takes O(log k) time.
4. Final Result: After checking all N vectors, the heap contains the k nearest neighbors. The distances will not be sorted, but you will have the correct set of k vectors. You can then extract them from the heap and sort them if a ranked list is required, which is a small O(k log k) operation.

This Stack Overflow answer provides a crystal-clear explanation of this exact process.

The overall complexity of this algorithm is O(N log k). Compared to the naive sorting approach of O(N log N), this is a massive improvement, as k is typically much, much smaller than N. Academic literature on k-NN confirms this is the best-known method for exact neighbor selection.

Your Task: Implementing a Bounded Heap in C

Now it's time to implement this in your project. You'll create a small set of C functions to manage a bounded max-heap.

First, you'll need a struct to hold the items in your heap, which will be a pair of {distance, index}. Then, define the heap structure itself.

// An item in our heap, storing the distance and the index of the vector
typedef struct {
    float distance;
    size_t index;
} Candidate;

// The bounded max-heap structure
typedef struct {
    size_t capacity; // k
    size_t count;    // current number of items in the heap
    Candidate* items; // array to store heap items
} BoundedHeap;

You will need to implement the following functions:

• BoundedHeap* heap_create(size_t capacity): Allocates memory for the heap structure and its internal items array.
• void heap_destroy(BoundedHeap* heap): Frees the allocated memory.
• void heap_push(BoundedHeap* heap, float distance, size_t index): The main function to add a candidate. This will contain the core bounded-heap logic.
• void sift_up(BoundedHeap* heap, size_t index): Helper to move an element up to restore the heap property (used when filling the initial k slots).
• void sift_down(BoundedHeap* heap, size_t index): Helper to move an element down to restore the heap property (used after replacing the root).

The image below visualizes these core operations, which are often called "swim up" (sift up) and "sink down" (sift down).

The logic inside heap_push is the most important part:

void heap_push(BoundedHeap* heap, float distance, size_t index) {
    if (heap->count < heap->capacity) {
        // Heap is not full yet, just add the new item and sift up
        heap->items[heap->count] = (Candidate){distance, index};
        sift_up(heap, heap->count);
        heap->count++;
    } else if (distance < heap->items[0].distance) {
        // Heap is full, but the new item is better than the worst in the heap
        // Replace the root and sift down
        heap->items[0] = (Candidate){distance, index};
        sift_down(heap, 0);
    }
    // Otherwise, the new item is worse than anything in the heap, so do nothing.
}

You will need to implement sift_up and sift_down using the 0-based index formulas for parent and children. A swap helper function will also be useful. This is a classic data structure implementation that will be a valuable addition to your low-level programming toolkit.

Integrate this heap into your brute-force k-NN search function. Instead of storing all results in a large array and sorting, create a heap of size k and push each distance calculation result into it. At the end of the loop, your heap will hold the k nearest neighbors.

Conclusion

In this lesson, you've moved from a naive O(N log N) selection algorithm to a highly efficient O(N log k) approach using a bounded max-heap. This is a fundamental technique in building high-performance search systems.

Key Takeaways:

• A naive sort-all approach to finding top k elements is inefficient for large datasets.
• A bounded max-heap is the ideal data structure for finding the k-smallest elements from a stream of data.
• The algorithm's efficiency comes from the O(1) rejection of non-candidate elements and the cheap O(log k) cost of updating the heap when a better candidate is found.
• Implementing a heap using an array provides excellent data locality and performance.

In our next lesson, we will analyze the performance of our fully optimized exact k-NN search. We will discover that even with these improvements, it can still be too slow for truly massive, high-dimensional datasets. This will set the stage for our exploration of Approximate Nearest Neighbor (ANN) search, where we'll begin to trade a small, controllable amount of accuracy for orders-of-magnitude improvements in search speed.