sklearn.neighbors.BallTree传递自定义度量函数时数据被修改的问题

嘿，这个问题我之前踩过坑！其实这是BallTree内部的数据类型转换机制在搞鬼——为了适配底层的浮点数运算优化，它会自动把你传入的整数数组转换成np.float64类型的数组，这就是你看到传入度量函数的数据和原始数据不一样的原因。

为啥会出现这种情况？ #

BallTree的核心逻辑是基于空间划分和浮点数距离计算实现的，sklearn的底层实现默认会把所有输入数据转成浮点类型（通常是np.float64），这样能保证后续近邻搜索的稳定性和效率。哪怕你传的是整数数组，它也会在初始化时悄悄完成这个转换，然后把转换后的浮点向量传给自定义度量函数。

拿你的示例代码来说，原始数据是整数类型：

print('Original data')
print(data)
# 输出：
# [[12 15  0]
#  [ 3  3  7]]

而传到度量函数里的是浮点版本：

print(x)
print(y)
# 输出：
# [12. 15.  0.]
# [ 3.  3.  7.]

值是完全一致的，只是数据类型从int变成了float而已。

怎么解决？ #

根据你的需求，有几种简单的处理方式：

1. 在度量函数里还原原始数据类型 #

如果你的度量逻辑必须依赖整数类型，直接把传入的浮点向量转回整数就行（因为值没变化，不会有精度损失）：

from sklearn.neighbors import BallTree
import numpy as np
np.random.seed(0)
data = np.random.randint(0, 20, size=(2, 3))

def metric(x, y):
    # 转换回整数类型，还原原始数据
    x_original = x.astype(np.int64)
    y_original = y.astype(np.int64)
    print('Data passed to metric (restored)')
    print(x_original)
    print(y_original)
    return 1

print('Original data')
print(data)
BallTree(data, metric=metric)

2. 提前把数据转成浮点类型 #

如果你不需要严格的整数类型，只是希望传入度量函数的数据和你输入BallTree的数据“看起来一致”，可以提前把原始数据转成浮点类型，这样BallTree就不会再做额外转换了：

from sklearn.neighbors import BallTree
import numpy as np
np.random.seed(0)
# 提前转换为浮点类型
data = np.random.randint(0, 20, size=(2, 3)).astype(np.float64)

def metric(x, y):
    print('Data passed to metric')
    print(x)
    print(y)
    return 1

print('Original data')
print(data)
BallTree(data, metric=metric)

3. 用闭包引用原始数据（复杂场景） #

如果你的度量逻辑需要依赖原始数据的其他额外属性（不仅仅是特征向量的值），可以把原始数据作为闭包变量传入度量函数，这样就能直接在函数里调用原始数据了：

from sklearn.neighbors import BallTree
import numpy as np
np.random.seed(0)
data = np.random.randint(0, 20, size=(2, 3))

def create_metric(original_data):
    def metric(x, y):
        # 这里可以直接使用原始数据做逻辑处理
        # 比如如果需要匹配原始数据的索引或其他字段
        print('Original data from closure')
        print(original_data)
        return 1
    return metric

print('Original data')
print(data)
# 传入原始数据创建度量函数
custom_metric = create_metric(data)
BallTree(data, metric=custom_metric)

验证效果 #

修改后再运行代码，你会发现传入度量函数的数据和原始数据的值完全一致，类型也能按需还原，完美解决数据被“修改”的问题。

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