Starting with Sklearns Nearest Neighbors output, how do I remove results where the record is it's own nearest neighbor?

Viewed 492

I am want to use sklearns NearestNeighbors() model to do some data analysis.

In my use case, I want to grab the N nearest neighbors and put it back into a pandas dataframe to evaluate the similarity of different records.

However, the results include the original record. In my case, that isn't useful. I want the nearest different records.

Example:

xtest = np.array([[1,1,1], [1,1,1], [1,.8,1] [.8,1,1]])
nn = NearestNeighbors(n_neighbors=2)
nn.fit(xtest)
distances, indices = nn.kneighbors(xtest)

returns:

(array([[0. , 0. ],
        [0. , 0. ],
        [0. , 0.2],
        [0. , 0.2]]),
 array([[0, 1],
        [0, 1],
        [2, 1],
        [3, 1]], dtype=int64))

In the above arrays the cells at indices (0,0), (1,1), (2, 0) and (3,0) are unimportant.

My goal is to manipulate this output so that I can create the following columns in pandas: "NearestNeighbor1" - the index of the nearest record other than itself "NearestNeighbor1_dist" -the distance of the nearest record other than itself even if the distance is zero.

"NearestNeighbor2" - the index of the next nearest record other than itself "NearestNeighbor2_dist" -the distance of the nearest record other than itself even if the distance is zero.

In the event of a tie, I don't care which record comes first (as long as it isn't itself).

2 Answers
from sklearn.neighbors import NearestNeighbors
import numpy as np
import pandas as pd

xtest = np.array([[1,1,1], [1,1,1], [1,.8,1], [.8,1,1]])
nn = NearestNeighbors(n_neighbors=3)
nn.fit(xtest)
distances, indices = nn.kneighbors(xtest)

df_ind = pd.DataFrame(data=indices)
df_ind = df_ind.apply(func=lambda x: [y for y in x if y != x.name], axis=1, result_type='expand')

df = pd.DataFrame({'NearestNeighbor1': df_ind.iloc[:, 0],
                   'NearestNeighbor1_dist': distances[:,1],
                   'NearestNeighbor2': df_ind.iloc[:, 1],
                   'NearestNeighbor2_dist': distances[:, 2]
                   })
print(df)

Output:

   NearestNeighbor1  NearestNeighbor1_dist  NearestNeighbor2  NearestNeighbor2_dist
0                 1                    0.0                 2                    0.2
1                 0                    0.0                 2                    0.2
2                 1                    0.2                 0                    0.2
3                 1                    0.2                 0                    0.2

This solution works for an arbitrary number of neighbors, although, I wonder if there is a more elegant solution using numpy.

N_NBRS = 4

nbrs = NearestNeighbors(n_neighbors=N_NBRS + 1, algorithm='brute')
nbrs.fit(X)

dist_n, ix_n = nbrs.kneighbors(X)
replacement = []
for row_idx, row in enumerate(ix_n):
    new_row = [val for val in row if val != row_idx]
    new_row = new_row[:N_NBRS] #Truncate in the event of many exact matches
    replacement.append(new_row)

ix_n2 = np.array(replacement)
dist_n2 = dist_n[:,1:]
results = X.copy()
for col_idx in range(N_NBRS):
    results[f'Neighbor{col_idx + 1}'] = ix_n2[:,col_idx]
    results[f'Neighbor{col_idx + 1}_dist'] = dist_n2[:,col_idx]
Related