如何用KNeighborsTransformer复现include_self=True的kneighbors_graph？

我的最终目标是用sklearn-ann包的转换器替代部分依赖kneighbors_graph的逻辑。sklearn-ann的所有方法都是兼容sklearn的转换器，但我要替换的逻辑使用了kneighbors_graph(mode="connectivity", include_self=True)，而我无法将转换器输出的不含自连接的距离结果转换成这种连接矩阵——并非所有转换器都支持含自连接的模式，但它们都能输出不含自连接的距离计算结果。

我可以从kneighbors_graph(mode="distance", include_self=True)（记为nn_with_self）复现目标连接矩阵，但无法从kneighbors_graph(mode="distance", include_self=False)（记为nn_without_self，与KNeighborsTransformer(mode="distance").fit_transform输出一致）复现。我发现nn_without_self是nn_with_self的超集，但不清楚后端筛选保留元素的规则。

我需要解决两个问题：

1. 如何从nn_without_self矩阵重建nn_with_self？
2. 全程如何基于稀疏矩阵操作，不转换为稠密矩阵？

我曾查阅sklearn后端代码，但类继承关系复杂，在GitHub上浏览多个文件后仍未理清逻辑。

代码示例 #

from sklearn.datasets import make_classification
from sklearn.neighbors import kneighbors_graph, KNeighborsTransformer
import numpy as np

X, _ = make_classification(n_samples=10, n_features=4, n_classes=2, n_clusters_per_class=1, random_state=0)
n_neighbors=3

# 生成K近邻矩阵
nn_with_self = kneighbors_graph(X, n_neighbors=n_neighbors, mode="distance", metric="euclidean", include_self=True,n_jobs=-1).todense()
nn_without_self = kneighbors_graph(X, n_neighbors=n_neighbors, mode="distance", metric="euclidean", include_self=False,n_jobs=-1).todense()
nn_from_transformer = KNeighborsTransformer(mode="distance", n_neighbors=n_neighbors, metric="euclidean", n_jobs=-1).fit_transform(X)

# 验证转换器输出与nn_without_self一致
print(np.all(nn_from_transformer == nn_without_self))
# 输出：True

print(np.all(nn_with_self == nn_without_self))
# 输出：False

# 检查矩阵对称性
print(np.allclose(nn_with_self,nn_with_self.T))
# 输出：False

print(np.allclose(nn_without_self,nn_without_self.T))
# 输出：False

实际矩阵输出 #

nn_with_self（含自连接） #

matrix([[0.        , 0.70550439, 0.        , 0.20463097, 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.51947869, 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.44145655],
        [0.        , 0.        , 0.        , 0.        , 0.50025504,
         0.        , 0.        , 0.        , 0.49481662, 0.        ],
        [0.20463097, 0.51947869, 0.        , 0.        , 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.50025504, 0.        , 0.        ,
         0.        , 0.        , 0.        , 0.34132965, 0.        ],
        [0.        , 0.88867318, 0.        , 0.        , 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.44956691],
        [0.        , 0.        , 1.10390699, 0.        , 1.52953542,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ,
         3.62670755, 0.        , 0.        , 0.        , 3.83571739],
        [0.        , 0.        , 0.49481662, 0.        , 0.34132965,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.44145655, 0.        , 0.        , 0.        ,
         0.44956691, 0.        , 0.        , 0.        , 0.        ]])

nn_without_self（不含自连接） #

matrix([[0.        , 0.70550439, 0.        , 0.20463097, 1.02852831,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.70550439, 0.        , 0.        , 0.51947869, 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.44145655],
        [0.        , 0.        , 0.        , 0.        , 0.50025504,
         0.        , 1.10390699, 0.        , 0.49481662, 0.        ],
        [0.20463097, 0.51947869, 0.        , 0.        , 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.95611187],
        [1.02852831, 0.        , 0.50025504, 0.        , 0.        ,
         0.        , 0.        , 0.        , 0.34132965, 0.        ],
        [0.        , 0.88867318, 0.        , 1.40547465, 0.        ,
         0.        , 0.        , 0.        , 0.        , 0.44956691],
        [0.        , 0.        , 1.10390699, 0.        , 1.52953542,
         0.        , 0.        , 0.        , 1.59848513, 0.        ],
        [0.        , 4.1280709 , 0.        , 0.        , 0.        ,
         3.62670755, 0.        , 0.        , 0.        , 3.83571739],
        [1.36553076, 0.        , 0.49481662, 0.        , 0.34132965,
         0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.44145655, 0.        , 0.95611187, 0.        ,
         0.44956691, 0.        , 0.        , 0.        , 0.        ]])

核心逻辑 #

当include_self=True时，kneighbors_graph会将每个样本自身作为第1个近邻（距离为0），再额外选取n_neighbors-1个其他样本作为近邻；而include_self=False时，直接选取n_neighbors个其他样本作为近邻。因此，从nn_without_self重建nn_with_self的关键是：对每个样本，保留其距离最小的n_neighbors-1个非自身样本，再添加自身的0距离连接。

稀疏矩阵实现代码 #

import scipy.sparse as sp

def reconstruct_nn_with_self(nn_without_self, n_neighbors):
    # 提取稀疏矩阵的非零元素行、列、值
    rows, cols, vals = sp.find(nn_without_self)
    
    filtered_rows = []
    filtered_cols = []
    filtered_vals = []
    
    # 遍历每个样本，筛选符合要求的近邻
    for i in range(nn_without_self.shape[0]):
        # 获取当前样本的所有非零近邻信息
        mask = rows == i
        row_cols = cols[mask]
        row_vals = vals[mask]
        
        # 按距离从小到大排序，取前n_neighbors-1个
        sorted_indices = np.argsort(row_vals)
        selected_indices = sorted_indices[:n_neighbors-1]
        
        # 添加筛选后的近邻
        filtered_rows.extend([i]*len(selected_indices))
        filtered_cols.extend(row_cols[selected_indices])
        filtered_vals.extend(row_vals[selected_indices])
        
        # 添加自身的0距离连接
        filtered_rows.append(i)
        filtered_cols.append(i)
        filtered_vals.append(0.0)
    
    # 重建稀疏矩阵
    return sp.csr_matrix((filtered_vals, (filtered_rows, filtered_cols)), 
                         shape=nn_without_self.shape)

验证代码 #

# 使用稀疏矩阵版本的nn_without_self测试
nn_without_self_sparse = kneighbors_graph(X, n_neighbors=n_neighbors, mode="distance", metric="euclidean", include_self=False,n_jobs=-1)
reconstructed_matrix = reconstruct_nn_with_self(nn_without_self_sparse, n_neighbors)

# 对比原nn_with_self
print(np.allclose(reconstructed_matrix.todense(), nn_with_self))
# 输出：True

说明 #

• 全程基于稀疏矩阵操作：通过sp.find提取非零元素，分组筛选后用sp.csr_matrix重建，未转换为稠密矩阵，适合处理大规模数据集。
• 若需要转换为连接矩阵（mode="connectivity"），只需将重建后的矩阵中所有非零值替换为1即可：

  connectivity_matrix = reconstructed_matrix.copy()
  connectivity_matrix.data = np.ones_like(connectivity_matrix.data)
  

内容的提问来源于stack exchange，提问作者O.rka