fork download
  1. import numpy as np
  2. from scipy.spatial.distance import cdist
  3. from sklearn.datasets import make_classification
  4.  
  5.  
  6. X, _ = make_classification(n_samples=20, n_features=4, n_classes=2, n_clusters_per_class=1, random_state=0)
  7. A, B = X[:10, :], X[10:, :]
  8. A.shape, B.shape
  9. dist = cdist(A, B)
  10. dist.shape
  11.  
  12. n_neighbors = 3
  13. inds = np.argsort(dist, axis=1)[:, :n_neighbors]
  14. dists = np.sort(dist, axis=1)[:, :n_neighbors]
  15.  
  16. print(inds, dists)
  17.  
Success #stdin #stdout 0.67s 75200KB
stdin
Standard input is empty
stdout
[[4 3 2]
 [7 8 0]
 [3 1 9]
 [1 9 3]
 [5 9 1]
 [8 7 0]
 [9 1 3]
 [4 3 1]
 [6 8 7]
 [4 3 1]] [[0.14806806 1.07924939 1.15818197]
 [0.48883274 0.6173809  0.92404217]
 [0.03744977 0.34636149 0.51052857]
 [0.05915844 0.2164162  0.33114403]
 [1.23264576 2.69432644 2.86056054]
 [0.35174069 0.75745427 1.35777454]
 [0.00825783 0.17480722 0.55386436]
 [0.45962337 0.50575253 0.88522623]
 [0.32848613 1.62949565 2.65980844]
 [0.39857567 0.80271613 1.15115862]]