Distributed Computing
Large Scale Data Processing
Parallel Algorithms
K-Way Merge
Data Management Systems

Find the largest k numbers in k arrays stored across k machines

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

When data is spread across many machines, shipping everything to one coordinator is usually the wrong answer. To find the global top k values from k arrays on k machines, the main goal is to minimize network traffic while still guaranteeing the correct result.

Core Idea

A coordinator does not need every element from every machine. It only needs enough candidates to reconstruct the final answer.

The standard pattern is:

  1. Each machine computes its local top k.
  2. Each machine sends only those k values to the coordinator.
  3. The coordinator merges the k local result sets and keeps the best global k.

If machine i stores n_i values, the local work is O(n_i log k) with a min-heap of size k. Network traffic is limited to at most k values per machine, which is far smaller than shipping full arrays.

Local Top k with a Min-Heap

Each worker scans its own array once. A min-heap is a good fit because the smallest candidate is always at the top. If a new number is larger than the smallest heap element, it belongs in the local top k.

python
1import heapq
2
3def local_top_k(values, k):
4    heap = []
5    for value in values:
6        if len(heap) < k:
7            heapq.heappush(heap, value)
8        elif value > heap[0]:
9            heapq.heapreplace(heap, value)
10    return sorted(heap, reverse=True)

This uses only O(k) extra memory regardless of array size.

Coordinator Merge

Once each machine sends its local result, the coordinator can apply the same idea again. The input is now much smaller: at most k * k values if there are k machines.

python
1def global_top_k(machine_results, k):
2    heap = []
3    for partial in machine_results:
4        for value in partial:
5            if len(heap) < k:
6                heapq.heappush(heap, value)
7            elif value > heap[0]:
8                heapq.heapreplace(heap, value)
9    return sorted(heap, reverse=True)
10
11
12workers = [
13    [4, 21, 7, 100, 33, 15],
14    [8, 9, 101, 3, 88, 60],
15    [5, 13, 55, 89, 120, 1],
16]
17
18k = 3
19partials = [local_top_k(worker, k) for worker in workers]
20answer = global_top_k(partials, k)
21
22print(partials)
23print(answer)

Output:

text
[[100, 33, 21], [101, 88, 60], [120, 89, 55]]
[120, 101, 100]

Why This Is Correct

Suppose a machine keeps only its local top k. Could the true global top k contain an element from that machine that was discarded locally? No. If an element is not in the local top k, then that same machine already has at least k larger elements. That discarded element cannot survive into the global top k.

That argument is the reason the algorithm is exact, not approximate.

Variants and Optimizations

If local arrays are already sorted in descending order, workers can return the first k elements immediately. That reduces local CPU cost to O(k) after storage access.

If communication is expensive, a tree-style reduction can help. Instead of one central coordinator receiving all results, pairs of machines merge partial top k sets, then forward only the merged top k upward. The communication volume per step stays small, and the reduction can scale better.

If machines are unreliable, you may also want to attach machine identifiers and use timeouts or retry logic. The algorithmic idea stays the same.

Common Pitfalls

The most common mistake is sorting every full array and sending all results over the network. That works, but it destroys the main advantage of distributed processing.

Another mistake is using a max-heap for fixed-size top k tracking. A max-heap makes removal awkward because you need fast access to the smallest retained item, not the largest.

A subtler issue is ties. If many machines contain the same value near the cutoff, decide whether you want exactly k values, k distinct values, or all values tied at the boundary. Those are different problems and need slightly different merge rules.

Summary

  • Compute local top k on each machine with a min-heap of size k.
  • Send only those candidates to the coordinator instead of full arrays.
  • Merge the partial results with the same min-heap technique.
  • The method is exact because discarded local values cannot enter the global top k.
  • Clarify tie handling and distinctness requirements before implementing the final merge.

Course illustration
Course illustration

All Rights Reserved.