How brute-force kneighbors finds the k nearest rows

1What brute-force kneighbors does, from start to finish

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

Brute force answers a nearest-neighbor question in the most direct way: it measures the distance from each query row to every training row and keeps the k smallest. It builds no index and skips no part of the data, so its answer is exact for any metric. The price is work that grows with the number of queries times the number of training rows, and most of this lesson is about how scikit-learn organizes that work so that it stays fast and fits in memory.

Every example on this page uses the same six training rows, each with two features. Row numbers start at 0, because NumPy counts positions from 0, and the result arrays of kneighbors contain these positions.

The model is fitted once, with k = 3 and brute force:

import numpy as np
from sklearn.neighbors import NearestNeighbors

X = np.array([[8, 6], [5, 2], [3, 0], [0, 0], [1, 8], [6, 9]], dtype=float)
nn = NearestNeighbors(n_neighbors=3, algorithm="brute").fit(X)
distances, indices = nn.kneighbors([[3, 3]])

What fit stores, and what kneighbors returns

For brute force, fit does almost no work. It checks the training rows and keeps them in the attribute _fit_X, a 2-D array with one row per training row, so N = 6 here. It records the method name "brute". And it settles the metric that every later step uses, in the attribute effective_metric_. The default metric is "minkowski" with p = 2, and fit renames it to "euclidean", because scikit-learn has faster code for that name. In the same way, p = 1 becomes "manhattan" and p = ∞ becomes "chebyshev". Every other metric keeps its name.

All of the search happens in kneighbors, which returns two arrays. Both have one row per query and k columns: indices[i, j] is the position of query i's j-th nearest training row, and distances[i, j] is the distance to it. Each row of results is sorted from nearest to farthest.

Read the two arrays

To check the result by hand, compute the squared distance from (3, 3) to each row, keep the three smallest, and take their square roots. The square root does not change which values are smallest, so the ranking can be done on squared distances.

1. Measure every row. The squared distances from (3, 3) to rows 0–5 are 34, 5, 9, 18, 29 and 45. For row 1: (5 − 3)² + (2 − 3)² = 4 + 1 = 5.
2. Keep the three smallest. 5 (row 1), 9 (row 2) and 18 (row 3).
3. Take square roots. √5 = 2.236, √9 = 3 and √18 = 4.243.
4. Read the arrays. indices is [[1, 2, 3]] and distances is [[2.236, 3.0, 4.243]]: one query row, three columns.

The six stages of every brute-force search

Under this one call, scikit-learn runs the same six stages for every input. Stage 2 picks one of two backends. The backends do stages 3 to 5 in different ways, but they return the same neighbors, so you can think of them as two engines behind one interface. Here's the whole path, with the section that explains each box:

Three facts from this picture carry through the rest of the lesson:

• The metric enters at one place only. Stage 3 fills a distance block, a table with one row per query and one column per training row, and no other stage measures anything. So stages 4 to 6 are the same for Euclidean, Manhattan, cosine, or a distance function that you wrote yourself, apart from one square root at the end of stage 4 that only Euclidean needs. That's why section 2 can treat every metric in one place.
• Both backends cut the work into chunks. The Python path takes some hundreds or thousands of queries at a time (section 3), and the compiled path cuts both sides into pieces of at most 256 rows (section 7). That's why the memory that a search needs stays small, however large the data is.
• The backends keep the k smallest in different ways. The Python path holds a whole row of distances for each query, so it can use argpartition on that row (section 4). The compiled path sees each query's distances a few hundred at a time, so it keeps a small heap per query instead (section 6).

The lesson follows the Python path first, in sections 2 to 5, because every line of it is NumPy code that you can run and print. Sections 6 and 7 then show what the compiled path does differently and when each backend runs.

No query: every row would find itself first

Stage 1 has one decision with consequences for the end of the search. You can call nn.kneighbors() with no argument. Then every training row becomes a query, and you get the neighbors of all six rows. Each row is at distance 0 from itself, so it would always be its own nearest neighbor. That result would be useless, and a classifier that used it would let every row vote for its own label.

The removal happens only when you pass no query. nn.kneighbors(X), with the training array passed in explicitly, treats it as ordinary queries, so each row does come back as its own nearest neighbor at distance 0. The user guide's first example shows exactly that.

So in the no-query case, scikit-learn asks for k + 1 neighbors, and at the end it removes each row's own index from its list. For the same reason, it accepts at most k = N − 1 here, because k + 1 rows must exist. A larger k raises a ValueError whose message says Expected n_neighbors < n_samples_fit. With a query, the message says <= instead, and k = N is allowed. The argument n_neighbors of kneighbors replaces the model's k for one call, and it must be a positive integer.

1. Ask for one more. k = 3, so the search keeps 4 per row.
2. Search. For row 0, the four nearest are rows 0, 5, 1 and 4, at distances 0, 3.606, 5 and 7.28.
3. Remove the row itself. Row 0 drops out, and its neighbors are rows 5, 1 and 4.

2Fill the distance block, with any metric

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

Stage 3 is the only stage that measures anything, so it's the stage where the metric matters. For a group of M queries and N training rows it builds a distance block: an M × N table whose cell [i, j] is the distance from query i to training row j. Every later stage reads only this table.

Brute force needs nothing else from a metric, so it is the one search method that works with every metric that scikit-learn accepts. A tree needs a property such as the triangle inequality to skip parts of the data, and the earlier lesson shows why cosine distance doesn't have it. Brute force skips nothing, so it works with any function that returns a number for a pair of rows. What changes from one metric to another is how fast the block gets filled. A formula that becomes one matrix product is fast, compiled code that handles one pair at a time is slower, and a Python function that is called once per pair is the slowest by far.

Who fills the block, for each kind of metric

This table names the function that fills the block on each backend. "Not used" means that this metric is not in the compiled backend's list, so the Python path runs instead; section 7 gives the exact rule.

A Python function is slow because Python runs every call itself. For one query and six training rows that's 6 calls. For kneighbors() with no query, when all rows fit in one chunk, it's 21 calls: _pairwise_callable calls your function once for each of the 15 pairs of different rows plus the 6 rows against themselves, and copies each pair's result to the cell on the other side of the diagonal. For 100,000 queries against 100,000 rows, the count runs into billions.

Switch between three metrics on the same query to see that the search stays the same and only the block row changes:

The same query under Manhattan and cosine distance

Manhattan distance adds up the absolute differences of the features, like a walk along the lines of a grid instead of a straight line:

1. Measure one row. For row 1 = (5, 2): |5 − 3| + |2 − 3| = 2 + 1 = 3.
2. Fill the block row. All six rows give [8, 3, 3, 6, 7, 9].
3. Keep the three smallest. Rows 1 and 2 at 3, and row 3 at 6. Rows 1 and 2 tie, and section 4 explains why the order of tied rows is not promised.
4. Compare with Euclidean. The same three rows win, but at different distances: 3, 3 and 6 instead of 2.236, 3 and 4.243.

Cosine distance changes the answer completely, because it ignores how far a row is and looks only at its direction. It treats each row as an arrow from the origin (0, 0) and measures 1 − cos(angle between the arrows), which is 0 for arrows that point the same way and 1 for arrows at a right angle. The query (3, 3) points at 45°, and row 0 = (8, 6) points almost the same way, so row 0 is the nearest at 0.0101, although it is the second-farthest row by Euclidean distance. The cosine neighbors are rows 0, 5 and 1. Row 3 = (0, 0) has no direction at all, and scikit-learn gives it cosine distance 1.

Rank by a cheaper number, and convert only the k winners

Many metrics end with a step that doesn't change which distances are smaller. Euclidean distance ends with a square root, and Minkowski distance with exponent p ends with a p-th root. Both roots keep the order of numbers that are 0 or larger: if a < b, then √a < √b. So the search can rank the rows by the number before the root, which scikit-learn calls the reduced distance, and take roots only for the k rows that it keeps. Minkowski distance with p = 3 shows the idea on a metric other than Euclidean:

1. Reduced distances. Add the cubes of the absolute differences. For row 1: 2³ + 1³ = 9. All six rows give [152, 9, 27, 54, 133, 243].
2. Keep the three smallest. 9, 27 and 54, from rows 1, 2 and 3.
3. Take three cube roots. 2.0801, 3 and 3.7798, which is exactly what kneighbors returns for p = 3.

The two backends use the idea to different extents. The compiled path ranks on reduced distances for every metric: each metric class in scikit-learn has a function rdist that computes the reduced distance and a function that converts it back. For Manhattan and Chebyshev there's no root to skip, so their reduced distance is the distance itself. The Python path uses the idea only for Euclidean: it asks for squared distances and takes the square roots at the very end (section 4). For every other metric, its block holds real distances from the start.

The Euclidean shortcut: every squared distance from one matrix product

Euclidean is the default metric, and it gets the fastest code on both backends. The direct way to fill its block subtracts each training row from each query, squares the differences and adds them up. In Python that is a loop over M × N pairs. In NumPy it is one subtraction with broadcasting, which builds an array of M × N × D differences for D features, D times larger than the block itself.

The way around this is in the formula. For two features, write out the squared distance and expand each square with (a − b)² = a² − 2ab + b²:

(q₀ − x₀)² + (q₁ − x₁)² = (q₀² + q₁²) − 2 · (q₀x₀ + q₁x₁) + (x₀² + x₁²)

The first bracket depends only on the query, the last bracket only on the training row, and only the middle one involves both. With more features each bracket gets longer, but the shape stays the same:

‖q − x‖² = ‖q‖²−2 · q · x+‖x‖²

Here ‖q‖² is the sum of the squares of q's features, its squared length, and q · x is the dot product: multiply the matching features and add. The squared lengths need one number per row, computed once. All the dot products together form one matrix product, which NumPy hands to BLAS, a library of tuned machine code for linear algebra. So the whole block costs one matrix product plus two additions, and no M × N × D array is ever built. The earlier lesson's section 4 measures how much faster this is.

One cell, by the expansion

Take the queries to be the training rows themselves, as in nn.kneighbors(), and compute the squared distance from row 0 to row 1.

1. Squared lengths. Row 0 = (8, 6): 64 + 36 = 100. Row 1 = (5, 2): 25 + 4 = 29.
2. Dot product. 8 × 5 + 6 × 2 = 52.
3. Combine. 100 − 2 × 52 + 29 = 25.
4. Check directly. The differences are 3 and 4, and 3² + 4² = 25.

Three NumPy steps build the whole block

scikit-learn turns the expansion into array code in three steps. Each step uses one NumPy feature, and the Python notes below show each feature on a tiny array first.

1. The squared lengths. row_norms(X, squared=True) gives the same values as (X ** 2).sum(axis=1): square every element, then add along each row (axis=1 means "across the columns of each row"). For the six rows that is the list [100, 29, 9, 0, 65, 117], of shape (6,).
2. All the dot products at once. In scikit-learn's code, X holds the query rows, with shape (M, D), and Y holds the training rows, with shape (N, D). For nn.kneighbors() both are the same six rows. X @ Y.T is the matrix product of the queries with the transposed training rows, an (M, N) table whose cell [i, j] is the dot product of query i with training row j.
3. Add the squared lengths to the right cells. ‖q‖² must go to every cell in query q's row, and ‖x‖² to every cell in training row x's column. Broadcasting does both additions without a loop.

import numpy as np
X = np.array([[1, 2], [3, 4]])            # 2 queries, 2 features
Y = np.array([[5, 6], [7, 8], [9, 10]])   # 3 training rows
print(X @ Y.T)

[[17 23 29]
 [39 53 67]]

a = np.array([1, 2])
b = np.array([10, 20, 30])
print(a[:, None] + b[None, :])

[[11 21 31]
 [12 22 32]]

With these two features, scikit-learn's function _euclidean_distances writes the whole expansion in three lines:

distances = -2 * (X @ Y.T)     # the dot products, times −2
distances += XX                # XX = row_norms(X)[:, None], a column
distances += YY                # YY = row_norms(Y)[None, :], a row

Cosine distance uses a matrix product too. cosine_distances first divides every row by its length, so that each row has length 1. The dot product of two rows of length 1 is the cosine of the angle between them, so one matrix product gives every cosine at once, and the block is 1 minus that product.

Rounding can make a distance negative, so scikit-learn clips at 0

The expansion subtracts two large numbers that are almost equal when two rows are close to each other and far from the origin. A float64 value keeps about 16 significant digits, so the subtraction can leave a small rounding error where the true answer is 0, and the error can be negative. Take the row (1000001.3, 1000004.9, 1000004.9). On this machine, the expansion gives −0.0009765625 for its distance to itself. The square root of a negative number is nan ("not a number"), which would make every later comparison with it false.

So right after the three lines above, scikit-learn runs np.maximum(distances, 0, out=distances), which replaces every negative value with 0 in the same array. When the queries are the training rows, it also sets each row's distance to itself to exactly 0, because the true value is known (section 3 shows where). cosine_distances clips in the same way, to the range from 0 to 2, which is the range of 1 − cos.

3Chunks keep the distance block small

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

Section 2 built one block for all queries at once. That works for six rows, but not for real data. A table of float64 values needs 8 bytes per cell. For 100,000 queries against 100,000 training rows, that is 8 × 100,000 × 100,000 = 80,000,000,000 bytes, which is 80 GB: more memory than most machines have. The search only needs the k smallest values of each row, so the table never has to exist all at once.

The Python path therefore takes the queries a few rows at a time. A group of consecutive query rows is a chunk. For each chunk it computes the chunk's distances to all N training rows, reduces them to k columns at once (section 4), and only then moves on. The largest array alive is one chunk's block.

Chunking cannot change the answer. The k nearest training rows of one query depend only on that query's own row of the block, never on the other queries, so the rows can be computed in groups of any size and in any order. Only the speed and the memory change.

What happens to one chunk

The function pairwise_distances_chunked runs the loop. Before the loop, it does two things once. It computes how many rows fit in one chunk. And for "mahalanobis" without a matrix VI of your own, it computes that matrix from the training rows, so that every chunk measures in the same way; with a query, it raises an error instead and asks for VI in metric_params. Then each chunk goes through the same five steps:

1. Cut out the chunk. Take query rows start to stop, a slice of the queries.
2. Fill its block. Call pairwise_distances on the chunk and all N training rows, with the metric's function from section 2. The block has shape (chunk rows, N).
3. Fix the self cells. Only for a Euclidean metric with no query: set each row's distance to itself to exactly 0.
4. Reduce the block. Call the reduce function, which keeps the k smallest per row (section 4). It returns two arrays of shape (chunk rows, k), and a check confirms that they still have one row per query.
5. Hand out the result. Give the reduced chunk to the caller. Nothing refers to the block any more, so Python frees its memory before the next chunk's block is built.

Sections 3 to 5 follow these steps in order. The rest of this section is about steps 1 and 3: how big a chunk is, how the slices are made, and where the self cells are.

How many rows fit in one chunk

The setting working_memory, in MiB (1 MiB = 2²⁰ = 1,048,576 bytes), is the budget for one chunk's block. It is 1024 MiB unless you change it with sklearn.set_config or sklearn.config_context. One row of the block holds N float64 values, which is 8 × N bytes, so get_chunk_n_rows computes

rows per chunk = ⌊ working_memory × 2²⁰ / (8 × N) ⌋, at least 1, at most M

where ⌊ ⌋ means rounding down. It never goes below one row, even when one row alone exceeds the budget; then it warns. The budget covers the block only. A metric function also needs memory while it works, such as the differences of Manhattan distance, and that memory comes on top of the budget, as a comment in scikit-learn's source says.

1. Bytes per row. N = 100,000 training rows: 8 × 100,000 = 800,000 bytes.
2. Budget in bytes. 1024 × 1,048,576 = 1,073,741,824 bytes.
3. Divide and round down. 1,073,741,824 / 800,000 = 1342.18…, so 1342 rows per chunk, and each block takes 1,073,600,000 bytes.

The slices, and a generator that makes them one at a time

The function gen_batches(n, size) produces the chunks as slice objects. A slice is Python's name for a range of positions: slice(0, 2) means positions 0 and 1, and X[slice(0, 2)] is the same as X[0:2]. The last slice is shorter when n is not a multiple of size.

gen_batches and pairwise_distances_chunked are generator functions, and that is what keeps only one chunk in memory.

def batches(n, size):
    for start in range(0, n, size):
        print("computing", start)
        yield slice(start, min(start + size, n))

gen = batches(5, 2)        # nothing is computed yet
print(next(gen))           # runs up to the first yield
print(list(gen))           # runs to the end

computing 0
slice(0, 2, None)
computing 2
computing 4
[slice(2, 4, None), slice(4, 5, None)]

For the walkthrough, the lesson sets working_memory to 96 bytes, so that the six rows split into chunks of two. One row of distances to six training rows needs 48 bytes, and 96 // 48 = 2 (// divides and rounds down). Step through the three chunks:

Finding each row's own cell inside a chunk

Step 3 of the chunk loop sets each row's distance to itself to 0. It does that only when two things hold: the metric is Euclidean (the code checks that the metric's function is euclidean_distances), and the queries are the training rows themselves, as in kneighbors(). For other metrics, the block keeps whatever the metric function computed.

Inside a chunk, row i of the block is training row start + i, where start is the chunk's first row. So its own cell is in column start + i. scikit-learn sets those cells to 0 with one line, D_chunk.flat[start :: N + 1] = 0.

block = np.arange(12).reshape(2, 6)   # 2 rows, N = 6 columns
block.flat[2::7] = -1                 # positions 2 and 9
print(block)

[[ 0  1 -1  3  4  5]
 [ 6  7  8 -1 10 11]]

1. Write the cell's position. Row i, column start + i: position i × N + start + i.
2. Collect the terms. That is start + i × (N + 1), so the positions start at start and step by N + 1.
3. Apply it. Chunk 2–3 of the six rows: start = 2 and N + 1 = 7, so the positions are 2 and 9.

4Keep the k smallest: argpartition, then sort only those k

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

After stage 3, each chunk is a block with N columns per row, and stage 4 must turn every row into its k smallest values, in ascending order, together with the columns they came from. On the Python path, one function does this: _kneighbors_reduce_func, which pairwise_distances_chunked calls once per chunk (step 4 of the chunk loop in section 3). It only compares numbers and never looks at the metric's formula, so it is the same code for every metric. Its one metric-dependent line is at the end: for Euclidean, it takes square roots.

Sort everything, or partition first

The obvious way is to sort each row and take the first k. That puts all N values in order, although N − k of them are thrown away right after. Sorting N values takes roughly N × log₂ N comparisons, so for N = 100,000 it does about 17 times the work of one pass over the row. The order matters only among the k winners, which suggests a cheaper plan: first find which k columns hold the smallest values, in about one pass, and then sort just those k.

print(np.argsort([30, 10, 20, 5]))

[3 1 2 0]

rng = np.random.default_rng(3)
row = rng.random(10_000)
first = np.argpartition(row, 499)[:500]
print(np.array_equal(np.sort(row[first]), np.sort(row)[:500]))   # the 500 smallest?
print(bool(np.all(np.diff(row[first]) > 0)))                      # in sorted order?

True
False

So np.argpartition(dist, k - 1, axis=1)[:, :k] gives each row's k smallest columns. The argument is k − 1 because positions count from 0: the k-th smallest value belongs at position k − 1.

Which columns survive

Here is row 0 of the first chunk, with k + 1 = 4 because the queries are the training rows. Its squared distances to rows 0–5 are [0, 25, 61, 100, 53, 13].

1. Partition at position 3. np.argpartition(row, 3) gives the columns [0, 5, 1, 4, 2, 3]. Column 4 sits at position 3, because 53 is the fourth-smallest value, and the columns before it hold 0, 13 and 25.
2. Keep the first four. [0, 5, 1, 4], with the values [0, 13, 25, 53].
3. Sort only those four. np.argsort([0, 13, 25, 53]) is [0, 1, 2, 3], so the order stays [0, 5, 1, 4].
4. Take four square roots. [0, 3.606, 5, 7.28].

On these short rows NumPy happens to return the first k already sorted, but that is not promised, and the 10,000-value row in the note above shows a case where it fails. So scikit-learn always sorts the k survivors, and so should you.

Ties have no promised order either. Under Manhattan distance, rows 1 and 2 are both at 3 from the query (section 2). This run listed row 1 first, but NumPy doesn't promise which of two equal values comes first, so a list with row 2 first would be just as correct. With k = 1 and a tie for first place, only one of the two rows is returned, and NumPy doesn't promise which one.

Pairing each row with its own columns

Once argpartition has picked k columns per row, the function must read the k values back, sort them, and put the columns in the same order. Each row has its own set of columns, so the code needs a way to say "row i, at row i's columns" for all rows at once.

dist = np.array([[5, 7, 9], [4, 6, 8]])
cols = np.array([[2, 0], [1, 2]])
rows = np.arange(2)[:, None]          # [[0], [1]]
print(dist[rows, cols])

[[9 5]
 [6 8]]

To read the k values back, scikit-learn writes dist[sample_range, neigh_ind]. Here neigh_ind has shape (chunk rows, k), and sample_range = np.arange(chunk rows)[:, None] is the column [[0], [1], …]. The whole reduce function for Euclidean distance is five lines:

sample_range = np.arange(dist.shape[0])[:, None]
neigh_ind = np.argpartition(dist, n_neighbors - 1, axis=1)
neigh_ind = neigh_ind[:, :n_neighbors]
neigh_ind = neigh_ind[sample_range, np.argsort(dist[sample_range, neigh_ind])]
result = np.sqrt(dist[sample_range, neigh_ind]), neigh_ind

In words, line by line:

1. Number the rows. sample_range is the column [[0], [1], …] with one entry per row of the chunk.
2. Partition every row. argpartition at position k − 1 moves each row's k smallest columns to the front, in no promised order.
3. Keep the front. [:, :n_neighbors] keeps the first k columns of every row.
4. Sort the survivors. dist[sample_range, neigh_ind] reads each row's k values, np.argsort returns the order that sorts them, and indexing neigh_ind with that order puts the columns in ascending order of distance.
5. Read the values and finish. The block holds squared distances for Euclidean, so the last line reads the sorted values once more and takes their square roots.

Step through one row at a time:

Square roots only for the survivors, and only for Euclidean

The Python path asks euclidean_distances for squared distances (squared=True) and takes square roots only in the last line above, on k values per row instead of N. The ranking stays correct, because for numbers that are 0 or larger, a smaller value always has a smaller square root. This is the reduced-distance idea of section 2, used for one metric.

For every other metric, the block already holds real distances, so the last line reads the values without a square root: result = dist[sample_range, neigh_ind], neigh_ind. And with kneighbors(..., return_distance=False), the function returns only neigh_ind, because no step after it needs the values.

5Join the chunks, then remove each row from its own list

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

After the chunk loop, the Python path holds a list with one reduced result per chunk, collected by list(pairwise_distances_chunked(...)) in kneighbors. Each item is a pair of a distance block and an index block, each of shape (chunk rows, k). Stage 5 joins them into one pair of arrays with one row per query. The compiled path skips stage 5, because it writes every query's result into one array from the start. Stage 6 runs on both backends, but only when the call had no query.

Stack the chunks

pairs = [("d0", "i0"), ("d1", "i1"), ("d2", "i2")]
dists, inds = zip(*pairs)
print(dists)
print(inds)
print(np.vstack([np.array([[1, 2]]), np.array([[3, 4], [5, 6]])]))

('d0', 'd1', 'd2')
('i0', 'i1', 'i2')
[[1 2]
 [3 4]
 [5 6]]

So neigh_dist, neigh_ind = zip(*chunked_results) splits the list into the distance blocks and the index blocks, and np.vstack stacks each group. The chunks come in query order, so the stacked rows do too.

1. Collect. Three chunks, each a pair of arrays of shape (2, 4).
2. Split. neigh_dist, neigh_ind = zip(*chunked_results) gives three distance blocks and three index blocks.
3. Stack. np.vstack(neigh_ind) has shape (6, 4): 2 + 2 + 2 rows of 4 columns.

With return_distance=False, each chunk's result is only an index block, so kneighbors calls np.vstack(chunked_results) directly, without zip.

Remove each row's own index with a mask

ind = np.array([[0, 2, 1], [1, 0, 2]])
mask = ind != np.arange(2)[:, None]
print(mask)
print(ind[mask])
print(ind[mask].reshape(2, 2))

[[False  True  True]
 [False  True  True]]
[2 1 0 2]
[[2 1]
 [0 2]]

scikit-learn compares the stacked index array with the column of row numbers, neigh_ind != np.arange(6)[:, None], so True marks another row and False marks the row itself. It keeps the True entries, reshapes them into rows of k, and applies the same mask to the distance array, so each distance stays with its index. The reshape works only because each row has exactly k True entries, and the corner case below is about the one situation where that fails.

1. Stacked indices of row 2. [2, 1, 3, 0].
2. Mask of row 2. [False, True, True, True], because only the first entry equals 2.
3. Keep and reshape. Row 2's neighbors are [1, 3, 0].

The corner case: duplicates can push a row out of its own list

Each list holds k + 1 entries. When a row has more than k identical copies, all of them are at distance 0, and the list may fill with copies before the row's own index gets a place. Then the row's mask is all True, the row has k + 1 entries to keep instead of k, and the reshape would fail.

This happens on the compiled path. Its heap rejects a value that ties with the worst value kept (section 6), so among tied rows the heap keeps the earlier ones. Take five rows, where rows 0, 1 and 2 are all (0, 0), with k = 1, so each list holds 2 entries:

1. Raw lists. Rows 0, 1 and 2 all get [0, 1]: two zeros, from the first two copies. Row 2's own index is missing.
2. Mask. Row 2's mask is [True, True]. np.all(mask, axis=1) finds exactly this row.
3. Fix. sample_mask[:, 0][dup_gr_nbrs] = False drops the first entry of such rows instead. Row 2 keeps [1], which is a copy of itself at distance 0, and every row again has k entries.

After the mask, kneighbors returns the two arrays, each of shape (N, k), and the search is complete. Everything from stage 3 on was the Python path. The next two sections show how the compiled path does stages 3 to 5 and when it runs instead.

6The compiled path keeps a heap per query instead

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

Sections 3 to 5 followed the Python path. For float64 or float32 data and most metrics, scikit-learn doesn't run that path at all. It runs a compiled backend called ArgKmin ("arg-k-min": the arguments, meaning column positions, of the k minimum values). ArgKmin is written in Cython, a language that turns Python-like code into C, so its loops run as machine code without the Python interpreter. Section 7 says exactly when it runs. This section explains how it keeps the k smallest, which is where it differs most from the Python path.

ArgKmin cuts the training rows into chunks too, so the distances for one query arrive a few hundred at a time, never as one complete row. argpartition needs the whole row at once, so it doesn't fit here. The backend needs a structure that keeps the k best candidates seen so far and updates them each time that a new distance arrives.

Keep the k best so far

For each new distance, there's one question: does it beat the worst of the k candidates kept so far? If it doesn't, the distance is thrown away. If it does, it replaces the worst one. So the structure must give the worst kept value at once and must replace it cheaply. Two candidates for that structure show why scikit-learn picked a heap:

• A sorted list of k values gives the worst value at once, because it's at the end. But keeping the list sorted after a replacement means shifting up to k entries to make room for the new one.
• A max-heap of k values also gives the worst value at once, because it's always in slot 0. A replacement needs at most ⌊log₂ k⌋ swaps: 2 for k = 5 and 6 for k = 100.

Both reject a value with one comparison. Once the heap holds good candidates, most of the N distances are rejected, so that one comparison is most of the work, and the cheaper replacement settles the choice for the heap.

A max-heap of k slots

A max-heap here is a list of k slots with one rule: every slot holds a value at least as large as the values of its children, where the children of slot s are slots 2s + 1 and 2s + 2. Slot 0 is the root, and slots 1 and 2 are its children. Because of the rule, slot 0 always holds the largest value in the heap, which is the worst of the k candidates. Nothing else is sorted. The tree exists only in these index formulas, so the heap is an ordinary block of k numbers in memory. Every slot starts at DBL_MAX, the largest float64, shown as ∞, so the first k real distances always enter.

Offering a value v to the heap (scikit-learn's heap_push) has two cases:

1. Reject. If v is at least as large as slot 0, it cannot be among the k smallest. That takes one comparison, and nothing changes. A tie with slot 0 is rejected too.
2. Replace and sift down. Otherwise v takes slot 0, and the old largest value is gone. Then v moves down: while a child of its slot holds a larger value, the larger child moves up into its slot, and v continues from that child's slot. It stops when no child is larger.

The row index travels with every value, so each slot holds both a value and the row that it belongs to. Step through the query (3, 3) and the six rows with k = 3, which is exactly the order in which ArgKmin offers them:

After the last training row

When every training row has been offered, each query's heap holds its k nearest rows, but in heap order and as reduced distances. Four steps turn the heaps into the result:

1. Merge, when threads split the training rows. When several threads each saw a part of the training rows, each thread has its own heaps for the same queries. The backend offers every value of those heaps to the main heap with the same heap_push, so the main heap ends with the k smallest of all of them.
2. Sort each heap. simultaneous_sort sorts each heap's values in ascending order and moves the row indices along with them, so each index stays with its value.
3. Turn reduced distances into distances. The heaps hold reduced distances (section 2). When you asked for distances, the backend clips each value at 0 and converts it: a square root for Euclidean, a p-th root for Minkowski, and nothing for Manhattan.
4. Return. The heaps already are the rows of two arrays of shape (queries, k), so nothing has to be stacked. Only stage 6 still runs in kneighbors: when there was no query, it removes each row from its own list.

Why a heap and not argpartition

Both keep the k smallest values, and for the same distances they keep the same rows, up to the order of ties. The difference is what they need to see. argpartition needs a query's whole row of N distances in memory, so the Python path stores (chunk rows × N) values. The heap needs one value at a time, so the compiled path stores k values per query plus one block of at most 256 × 256 distances per thread (section 7), however large N is.

7Which backend runs, and what the compiled path computes

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

So there are two backends for the same search: the compiled ArgKmin with heaps, and the Python path with chunks of queries and argpartition. In stage 2, kneighbors asks ArgKmin.is_usable_for(X, Y, metric) whether the compiled backend may run. Its conditions are about the numbers' type, the way that they are stored, and the metric, and the Python note below explains those words first.

a = np.array([[1, 2], [3, 4]])
print(a.dtype, a.flags.c_contiguous)
print(np.asfortranarray(a).flags.c_contiguous)
print(np.array([[1.0, 2.0]]).dtype)

int64 True
False
float64

First, one exception: when the queries and the training rows are both sparse and the metric's name contains "euclidean", is_usable_for says no at once, because the compiled code is too slow for that case when there are many features. Otherwise it returns True only when six conditions all hold:

1. the setting enable_cython_pairwise_dist is on (the default);
2. the queries are a dense NumPy array in C order or a valid CSR matrix;
3. the training rows are too;
4. both have the same dtype;
5. that dtype is float32 or float64;
6. the metric is in the compiled list, which leaves out a Python function as the metric, "mahalanobis", "hamming" and the boolean metrics, and which has no "cosine" at all.

The checks are joined by and, so Python stops at the first one that fails. Try the presets. Each one changes one thing from the default, and the last column shows what a real fit and kneighbors on that input did:

Integer rows are the case that surprises people most. fit keeps integer training rows as integers, so kneighbors on them fails condition 5 and runs the Python path, even with the Euclidean metric. Converting the data with X.astype(float) before fit moves the search to the compiled path.

Inside the compiled path: two classes

When is_usable_for returns True, ArgKmin.compute picks the compiled class by dtype (ArgKmin64 for float64, ArgKmin32 for float32) and then by metric. Two classes do the work, and they differ only in stage 3, the way that they fill a block:

• EuclideanArgKmin64, for "euclidean" and "sqeuclidean", uses the expansion of section 2. It computes ‖x‖² for every training row and ‖q‖² for every query once, at the start. For each pair of chunks, one matrix product fills the middle terms, and then a loop adds the squared lengths, clips each value at 0, and offers it to its query's heap.
• The generic ArgKmin64, for every other compiled metric, has two nested loops over a pair of chunks, one over its queries and one over its training rows. For each pair of rows it calls the metric's rdist and offers the result to the heap at once, so it never builds a block at all.

Both classes cut both sides into chunks of at most 256 rows (the setting pairwise_dist_chunk_size). One query chunk and one training chunk need a block of at most 256 × 256 = 65,536 values, which is 524,288 bytes (512 KiB), small enough to stay in the processor's cache, the small fast memory on the chip. Reading from the cache is many times faster than reading from main memory, so working on one small block at a time keeps the processor busy instead of waiting for data.

gemm: the middle terms of a whole chunk pair in one call

For each chunk pair, EuclideanArgKmin64 gets the dot products from one call to gemm, the BLAS routine for a general matrix product. It computes

C = α · A · B + β · C

where α and β are plain numbers. scikit-learn passes the query chunk as A, the training chunk as B with the flag "read it transposed", α = −2 and β = 0. So C becomes −2 · q · x for every pair in the chunk pair: the middle term of the expansion, already multiplied by −2, with nothing left over from what C held before. A loop then adds ‖q‖² and ‖x‖² to each cell, clips it at 0, and offers it to the query's heap.

1. A and B. A is the query chunk [[3, 3]] and B is the six training rows.
2. One cell. For row 1 = (5, 2): q · x = 3 × 5 + 3 × 2 = 21, so C = −2 × 21 = −42.
3. The whole block. [−84, −42, −18, 0, −54, −90].
4. Add the norms. For row 1: 18 + (−42) + 29 = 5, the squared distance from section 1.

Threads, and what n_jobs does

A thread is one sequence of work that a processor core runs, and several threads run at the same time on several cores. ArgKmin starts its threads with OpenMP, a standard tool that compiled code uses to run work in parallel. By default it uses one thread per physical CPU core, unless the environment variable OMP_NUM_THREADS sets another number.

The threads share out one side of the problem. The backend gives each thread its own query chunks when there are more queries than training rows, or more than four chunks of 256 queries per thread; otherwise it gives each thread its own training chunks. With one query, as in this lesson, it shares out the training rows. Each thread then fills its own heaps, and the merge at the end of section 6 combines them. While ArgKmin runs, it limits BLAS to one thread, because each of its own threads calls BLAS, and a BLAS call inside each of them must not start more threads.

The estimator's n_jobs setting doesn't change any of this, because ArgKmin.compute doesn't receive it. It matters only on the Python path: kneighbors passes it to pairwise_distances, which then splits the training rows of each chunk's block among that many threads.

8Walk through every line scikit-learn runs

This section follows the real source of two calls on the fitted model from section 1, recorded with a tracer on scikit-learn 1.8.0. Line numbers are the installed files' numbers. Lines that ran are shown normally, and lines that did not run are dimmed.

• nn.kneighbors([[3, 3]]) with the default settings takes the compiled ArgKmin path. Its algorithm runs in Cython, which a tracer cannot see, so those steps show the template files that scikit-learn ships (.pyx.tp), and their values come from the finished results.
• nn.kneighbors(), with no query, runs inside config_context(enable_cython_pairwise_dist=False, working_memory=96 / 2**20). That turns the compiled backend off and makes chunks of two rows, so the Python path runs every branch of sections 3–5. scikit-learn documents the switch as a setting for benchmarking and testing; the lesson uses it only to show the Python path on the same data.

Start with the call tree to see the whole path, then read the steps in order. Each step names its subgoal, explains the code for any input first, and then says what happened in this run.

The steps

9Brute-force kneighbors from scratch

1. 1 Prepare the call
2. 2 Choose a backend
3. 3 Fill the distance blocks
4. 4 Keep the k smallest
5. 5 Join the pieces
6. 6 Remove each row from its own list

This implementation shows each step that the earlier sections took apart, in about 70 lines of NumPy code, not counting docstrings and comments. It follows the Python path, because every line of that path is NumPy you can read, and it skips stage 2, because it has only one backend. A second, slower function follows the compiled path for one query, with a heap written out in Python. It is written for reading, not speed. It handles only the Euclidean metric, and it leaves out input validation, sparse input and threads. It also computes the training rows' squared lengths once instead of once per chunk.

How the main function runs

Read kneighbors_brute first, because it calls every other function in the file. Its body is the six stages of section 1 without stage 2:

1. Prepare the call. When queries is None, it uses the training rows as the queries and searches for k + 1. It raises a ValueError when the search needs more rows than the training rows have.
2. Size the chunks. It computes the rows per chunk from the memory budget, with the formula of section 3, and the squared lengths of the training rows, once for all chunks.
3. Fill each chunk's block. In the loop over chunks, squared_distances computes the block with the expansion and the clip at 0. With no query, the next line sets each row's distance to itself to 0.
4. Keep the k smallest. k_smallest partitions each row, sorts the survivors, and returns their squared distances and columns. The loop takes the square roots and appends both arrays to two lists.
5. Join the pieces. After the loop, np.vstack stacks each list into one array with one row per query.
6. Remove each row from its own list. With no query, drop_self applies the mask, including the fix for rows with many exact copies.

heap_push and kneighbors_heap stand apart at the end of the file. They do stage 4 the compiled way, for one query: every slot starts at infinity, every squared distance is offered to the heap, and the k survivors are sorted at the end.

Each function has a docstring that says what it receives and returns, and comments mark the lines that repeat a rule from the lesson.

"""Brute-force k-nearest-neighbor search, written to show each step that scikit-learn's kneighbors takes.

The vectorized function ``kneighbors_brute`` follows scikit-learn's Python
fallback path; the loop function ``kneighbors_heap`` follows the compiled
ArgKmin path for one query. Both return what ``NearestNeighbors.kneighbors``
returns, and the lesson's verification script checks them against it.

Stages
------
1. ``squared_distances``: squared Euclidean distances from a chunk of queries
   to every training row, with ``|q|^2 - 2 q.x + |x|^2``, clipped at 0.
2. ``kneighbors_brute``: split the queries into chunks that fit a memory
   budget, and set each row's distance to itself to exactly 0 when the
   queries are the training rows.
3. ``k_smallest``: keep the k smallest values in each row with
   ``np.argpartition``, then sort only those k.
4. ``drop_self``: when the queries are the training rows, remove each row
   from its own list of k + 1.
5. ``heap_push`` and ``kneighbors_heap``: the same search for one query with
   a size-k max-heap, which is how the compiled backend keeps the k smallest.

Execution behavior
------------------
Importing the module only defines the functions. Running it as a script
prints the lesson's six-row example.
"""

import numpy as np


def squared_distances(queries, train, train_sq_norms):
    """Return the squared Euclidean distance from every query row to every training row.

    The direct formula subtracts every query from every training row, which
    makes an array of shape (n_queries, n_train, n_features). The expansion
    ``|q|^2 - 2 q.x + |x|^2`` needs only one matrix product, and NumPy hands
    that product to BLAS. Rounding can leave values such as -1e-12 where the
    true value is 0, so the result is clipped at 0, as scikit-learn does.

    Parameters
    ----------
    queries : ndarray of shape (n_queries, n_features)
        The query rows of one chunk.
    train : ndarray of shape (n_train, n_features)
        All training rows.
    train_sq_norms : ndarray of shape (n_train,)
        ``(train ** 2).sum(axis=1)``, computed once by the caller and reused
        for every chunk.

    Returns
    -------
    ndarray of shape (n_queries, n_train)
        Entry [i, j] is the squared distance from query i to training row j.
    """
    query_sq_norms = (queries ** 2).sum(axis=1)

    # Same order of operations as scikit-learn's _euclidean_distances, so
    # rounding, and therefore the order of tied rows, matches it exactly.
    sq = -2 * (queries @ train.T)
    sq += query_sq_norms[:, None]      # column: one value per query row
    sq += train_sq_norms[None, :]      # row: one value per training row
    np.maximum(sq, 0, out=sq)
    return sq


def k_smallest(sq_distances, k):
    """Return the k smallest values of each row and their column positions, in ascending order.

    ``np.argpartition`` moves the k smallest values of each row into its
    first k positions without sorting the rest, which is cheaper than a full
    sort. Only those k values are then sorted.

    Parameters
    ----------
    sq_distances : ndarray of shape (n_rows, n_train)
        Squared distances of one chunk.
    k : int
        How many values to keep per row; 1 <= k <= n_train.

    Returns
    -------
    values : ndarray of shape (n_rows, k)
        The k smallest squared distances of each row, ascending.
    indices : ndarray of shape (n_rows, k)
        Their column positions, which are training-row indices.
    """
    rows = np.arange(len(sq_distances))[:, None]

    # Position k - 1 holds the k-th smallest value; positions 0 .. k - 1 hold
    # the k smallest, in no particular order.
    candidates = np.argpartition(sq_distances, k - 1, axis=1)[:, :k]

    # rows pairs each row number with that row's own k columns.
    order = np.argsort(sq_distances[rows, candidates], axis=1)
    indices = candidates[rows, order]
    return sq_distances[rows, indices], indices


def drop_self(distances, indices):
    """Remove each training row from its own list of k + 1 neighbors.

    Usually a row's own index is in its list, at distance 0. When the row has
    many exact duplicates, the duplicates can fill the list and leave the row
    out; then the first entry is dropped instead, so every row keeps k.

    Parameters
    ----------
    distances : ndarray of shape (n_rows, k + 1)
    indices : ndarray of shape (n_rows, k + 1)
        Row i lists the neighbors of training row i.

    Returns
    -------
    distances : ndarray of shape (n_rows, k)
    indices : ndarray of shape (n_rows, k)
    """
    n_rows, k_plus_one = indices.shape
    is_other_row = indices != np.arange(n_rows)[:, None]

    # Rows whose own index is missing: every entry is another row.
    self_missing = is_other_row.all(axis=1)
    is_other_row[self_missing, 0] = False

    shape = (n_rows, k_plus_one - 1)
    return distances[is_other_row].reshape(shape), indices[is_other_row].reshape(shape)


def kneighbors_brute(train, queries=None, k=5, working_memory_mib=1024):
    """Find the k nearest training rows of each query by measuring every distance.

    This follows scikit-learn's Python brute-force path: chunks of queries,
    squared distances by the expansion, argpartition per chunk, square roots
    of the kept values only, and, when ``queries`` is None, one extra
    neighbor that is removed at the end.

    Parameters
    ----------
    train : ndarray of shape (n_train, n_features)
        The training rows.
    queries : ndarray of shape (n_queries, n_features) or None, default=None
        The query rows. None uses the training rows as queries, and then no
        row counts as its own neighbor.
    k : int, default=5
        Number of neighbors per query.
    working_memory_mib : float, default=1024
        Memory budget for one chunk of squared distances, in MiB, as in
        scikit-learn's ``working_memory`` setting. A chunk always has at
        least one row.

    Returns
    -------
    distances : ndarray of shape (n_queries, k)
        Euclidean distances to the neighbors, ascending in each row.
    indices : ndarray of shape (n_queries, k)
        Row positions of the neighbors in ``train``.

    Raises
    ------
    ValueError
        When the search needs more rows than ``train`` has: k > n_train with
        queries, or k >= n_train without them.
    """
    query_is_train = queries is None
    if query_is_train:
        queries = train
        k_search = k + 1        # room for the row itself, removed at the end
    else:
        k_search = k
    if k_search > len(train):
        raise ValueError(f"need {k_search} neighbors but there are only {len(train)} training rows")

    # One row of the chunk holds len(train) float64 values of 8 bytes each.
    rows_per_chunk = max(1, int(working_memory_mib * 2**20 // (8 * len(train))))
    train_sq_norms = (train ** 2).sum(axis=1)

    distance_chunks, index_chunks = [], []
    for start in range(0, len(queries), rows_per_chunk):
        stop = min(start + rows_per_chunk, len(queries))
        sq = squared_distances(queries[start:stop], train, train_sq_norms)

        # Chunk row i is training row start + i: its distance to itself is 0.
        if query_is_train:
            chunk_rows = np.arange(stop - start)
            sq[chunk_rows, start + chunk_rows] = 0

        sq_kept, index_kept = k_smallest(sq, k_search)
        distance_chunks.append(np.sqrt(sq_kept))
        index_chunks.append(index_kept)

    distances, indices = np.vstack(distance_chunks), np.vstack(index_chunks)
    if query_is_train:
        distances, indices = drop_self(distances, indices)
    return distances, indices


def heap_push(values, indices, value, index):
    """Offer one value to a fixed-size max-heap stored in two lists.

    Slot 0 holds the largest kept value. A value that is not smaller than it
    is rejected after one comparison. Otherwise it replaces slot 0 and moves
    down: while a child of its slot holds a larger value, the larger child
    moves up. The children of slot s are slots 2s + 1 and 2s + 2. This is
    scikit-learn's ``heap_push`` from ``sklearn/utils/_heap.pyx``.

    Parameters
    ----------
    values : list of float
        Kept values in heap order; changed in place.
    indices : list of int
        The row index of each kept value; changed in place.
    value : float
        The new value.
    index : int
        Its row index.
    """
    if value >= values[0]:
        return

    size = len(values)
    slot = 0
    while True:
        left, right = 2 * slot + 1, 2 * slot + 2
        if left >= size:
            break
        has_right_child = right < size
        larger_child = right if has_right_child and values[right] > values[left] else left
        if values[larger_child] <= value:
            break
        values[slot], indices[slot] = values[larger_child], indices[larger_child]
        slot = larger_child
    values[slot], indices[slot] = value, index


def kneighbors_heap(train, query, k):
    """Find the k nearest training rows of one query with a heap, one distance at a time.

    This mirrors the compiled ArgKmin loop for a single query: every slot
    starts at infinity, every squared distance is offered to the heap, and
    the k survivors are sorted at the end. It computes each distance
    directly, without the expansion, to keep the loop short.

    Parameters
    ----------
    train : ndarray of shape (n_train, n_features)
    query : ndarray of shape (n_features,)
    k : int

    Returns
    -------
    distances : list of float
        Euclidean distances, ascending.
    indices : list of int
        Row positions in ``train``.
    """
    values, indices = [np.inf] * k, [-1] * k
    for row_index, row in enumerate(train):
        heap_push(values, indices, float(((query - row) ** 2).sum()), row_index)

    survivors = sorted(zip(values, indices))
    return [float(np.sqrt(value)) for value, _ in survivors], [index for _, index in survivors]


if __name__ == "__main__":
    train = np.array([[8, 6], [5, 2], [3, 0], [0, 0], [1, 8], [6, 9]], dtype=float)
    print(kneighbors_brute(train, np.array([[3.0, 3.0]]), k=3))
    print(kneighbors_heap(train, np.array([3.0, 3.0]), k=3))
    print(kneighbors_brute(train, k=3, working_memory_mib=96 / 2**20))

Two small choices differ from scikit-learn but give the same arrays. The self cells are set with pair indexing, sq[chunk_rows, start + chunk_rows] = 0, which names the cells (i, start + i) directly instead of computing their .flat positions. And heap_push picks the larger child in one line instead of three branches, which compares the same values.

Run it on the six rows

The three calls at the bottom of the file print the result for the query (3, 3), the heap version of the same search, and the neighbors of every training row with chunks of two rows:

(array([[2.23606798, 3.        , 4.24264069]]), array([[1, 2, 3]]))
([2.23606797749979, 3.0, 4.242640687119285], [1, 2, 3])
(array([[3.60555128, 5.        , 7.28010989],
       [2.82842712, 5.        , 5.38516481],
       [2.82842712, 3.        , 7.81024968],
       [3.        , 5.38516481, 8.06225775],
       [5.09901951, 7.21110255, 7.28010989],
       [3.60555128, 5.09901951, 7.07106781]]), array([[5, 1, 4],
       [2, 0, 3],
       [1, 3, 0],
       [2, 1, 4],
       [5, 1, 0],
       [0, 4, 1]]))

These are the arrays from sections 1 and 5. The lesson's verification script also compares the functions with NearestNeighbors.kneighbors on 300 random problems of different sizes, k and memory budgets, with a query and without one, and on the duplicates example of section 5. Every result was identical.

Change the metric

Section 2 said that only stage 3 depends on the metric, and the code shows it. To search with Manhattan distance, only two lines of kneighbors_brute change: the line that fills the block, and the square root, which Manhattan doesn't need. This function fills a Manhattan block:

def manhattan_block(queries, train):
    differences = queries[:, None, :] - train[None, :, :]
    return np.abs(differences).sum(axis=2)

queries[:, None, :] has shape (M, 1, D) and train[None, :, :] has shape (1, N, D), so broadcasting gives an (M, N, D) array of differences, one per query, training row and feature. np.abs takes their absolute values, and sum(axis=2) adds them over the features, which leaves the (M, N) block. This is the large array that the Euclidean expansion avoids, but the chunks keep M small, so it stays within a few times the block's own memory. In the loop, sq = squared_distances(...) becomes block = manhattan_block(queries[start:stop], train), and np.sqrt(sq_kept) becomes just the kept values. A row's Manhattan distance to itself is exactly 0 without any rounding, so the self-cell line can stay or go. The verification script runs this function through k_smallest and drop_self and compares the result with NearestNeighbors(metric="manhattan").

Fill in the missing step

rows = np.arange(len(sq_distances))[:, None]
candidates = np.argpartition(sq_distances, k - 1, axis=1)[:, :k]
# missing: the order that sorts each row's k values
indices = candidates[rows, order]
return sq_distances[rows, indices], indices

n_rows, k_plus_one = indices.shape
is_other_row = indices != np.arange(n_rows)[:, None]
# missing: find the rows whose own index is not in their list
# missing: for those rows, drop the first entry instead
shape = (n_rows, k_plus_one - 1)
return distances[is_other_row].reshape(shape), indices[is_other_row].reshape(shape)

Find the bug

Each of these questions describes a plausible change to the code above. Find the smallest input on which the change gives a wrong answer or an error.

10Check yourself without looking back

Try these from memory, without scrolling back. Each answer names the section that teaches it.

11Review cards

Revisit these after a day, a week, and a month. Each card makes sense without the page open.

Why does scikit-learn's kneighbors() with no query ask for k + 1 neighbors?::Every training row is then a query and is at distance 0 from itself, so one extra neighbor is found and each row's own index is removed at the end.
In a brute-force search with kneighbors, which stage depends on the metric?::Only stage 3, which fills the distance block. The later stages compare and move numbers, apart from the square root that Euclidean needs at the end.
What is a reduced distance, and why does brute force rank by it?::A cheaper number with the same order as the distance, such as the squared Euclidean distance or the Minkowski sum before its p-th root. Ranking by it gives the same neighbors, and only the k kept values need the root.
Write the expansion that brute-force Euclidean KNN uses for ‖q − x‖².::‖q‖² − 2·q·x + ‖x‖². Only the dot products involve both rows, and all of them come from one matrix product.
What does a[:, None] + b[None, :] produce for a of shape (M,) and b of shape (N,)?::An (M, N) array whose cell [i, j] is a[i] + b[j], by broadcasting.
Why does scikit-learn clip squared distances at 0 after the expansion?::Rounding can leave a small negative number where the true value is 0, and its square root would be nan.
How many query rows fit in one chunk when scikit-learn runs brute force in Python?::working_memory (MiB) × 2²⁰ ÷ (8 × number of training rows), rounded down, at least 1, at most the number of queries.
In a chunk starting at row s of a set of N training rows, where are the self cells in D_chunk.flat?::At positions s, s + (N + 1), s + 2(N + 1), …, which D_chunk.flat[s :: N + 1] = 0 sets to zero.
What does np.argpartition(row, k - 1)[:k] guarantee, and what does it not?::It gives the columns of the k smallest values, but in no promised order, so an argsort of those k must follow.
When does the duplicates fix in kneighbors() drop the first entry of a row's list?::When the row's own index is missing from its k + 1 entries, because more than k exact copies of the row filled the list.
What does a size-k max-heap keep in slot 0, and when does heap_push reject a value?::Slot 0 keeps the largest of the k kept values. heap_push rejects a value after one comparison when the value is at least as large as slot 0.
Name three inputs that make brute-force kneighbors skip the compiled ArgKmin backend.::Integer training rows, a Python function or 'mahalanobis' or 'hamming' as the metric, and the setting enable_cython_pairwise_dist=False.
Which gemm call does EuclideanArgKmin make, in terms of α and β?::C = α·A·Bᵀ + β·C with A the query chunk, B the training chunk, α = −2 and β = 0, which gives −2·q·x for every pair.

Checked on 2026-11-04 with scikit-learn 1.8.0 and NumPy 2.4.4: 1521 checks passed.

Sources

• scikit-learn 1.8.0 source, read from the installed package (BSD-3-Clause license): sklearn/neighbors/_base.py (kneighbors, _kneighbors_reduce_func); sklearn/metrics/pairwise.py (pairwise_distances_chunked, pairwise_distances, euclidean_distances, _euclidean_distances); sklearn/utils/_chunking.py; sklearn/metrics/_pairwise_distances_reduction/ (_dispatcher.py, _base.pyx.tp, _argkmin.pyx.tp, _middle_term_computer.pyx.tp); sklearn/utils/_heap.pyx; sklearn/_config.py.
• scikit-learn 1.8 user guide, Nearest Neighbors: the brute-force section and the kneighbors(X) example.
• scikit-learn 1.8, set_config: working_memory, pairwise_dist_chunk_size and enable_cython_pairwise_dist.
• NumPy 2.4, numpy.argpartition: "The order of all elements in the partitions is undefined."
• NumPy 2.4, Broadcasting: the general broadcasting rules.
• How scikit-learn KNN chooses brute force, KDTree, or BallTree: the earlier lesson, for the cost of brute force and how much BLAS speeds it up.