在Databricks PySpark Notebook中KD-Tree与Ball-Tree算法结果不符预期

Your output is correct given the default Euclidean distance metric used by BallTree and KDTree. Let's break down why:

Step 1: Calculate Euclidean Distances for Row 0 #

The first row of your matrix is [0, 1, 4, 2, 1]. Let's compute its Euclidean distance to every other row:

• Row 0 (self): Distance = 0
• Row 1:
  sqrt((0-1)² + (1-0)² + (4-2)² + (2-1)² + (1-1)²) = sqrt(1+1+4+1+0) = sqrt(7) ≈ 2.6458
• Row 3:
  sqrt((0-2)² + (1-1)² + (4-3)² + (2-0)² + (1-1)²) = sqrt(4+0+1+4+0) = sqrt(9) = 3
• Row 4:
  sqrt((0-1)² + (1-1)² + (4-1)² + (2-1)² + (1-0)²) = sqrt(1+0+9+1+1) = sqrt(12) ≈3.464
• Row 2:
  sqrt((0-4)² + (1-2)² + (4-0)² + (2-3)² + (1-1)²) = sqrt(16+1+16+1+0)=sqrt(34)≈5.830

The ordered distances are 0 < 2.6458 <3 <3.464 <5.830, which maps to indices [0,1,3,4,2]—exactly what your output shows.

Step 2: Did You Intend a Different Distance Metric? #

If you expected a different order (like [0,1,4,3,2]), you might have assumed a non-default distance metric. For example, using Manhattan distance:

• Row 0 to Row1: 1+1+2+1+0=5
• Row0 to Row4: 1+0+3+1+1=6
• Row0 to Row3: 2+0+1+2+0=5

With Manhattan distance, Row1 and Row3 have the same distance from Row0, so their order could vary. To use this metric, initialize the tree with metric='manhattan':

import numpy as np
from sklearn.neighbors import BallTree
matrix = np.array([
          [0, 1, 4, 2, 1],
          [1, 0, 2, 1, 1],
          [4, 2, 0, 3, 1],
          [2, 1, 3, 0, 1],
          [1, 1, 1, 1, 0]
      ])
tree = BallTree(matrix, metric='manhattan')
dist, ind = tree.query(matrix, k=5)
print(ind)

This might produce an order like [0,1,3,4,2] or [0,3,1,4,2] (since distances are tied), but not exactly your expected [0,1,4,3,2]—since Row3's Manhattan distance is still smaller than Row4's.

Conclusion #

Your original code is working as intended. The discrepancy comes from a mismatch between your expected distance metric and the default Euclidean metric used by scikit-learn's tree implementations.

内容的提问来源于stack exchange，提问作者Jonny