实现KNN关联矩阵时，如何结合使用np.fromiter与np.ndenumerate？

需求背景 #

手动实现将一组点的KNN图表示为关联矩阵：对n×n的亲和矩阵（元素代表点间距离/相似度）的每一行，枚举并排序后取前K个元素的索引。

原始代码及错误 #

原始尝试代码 #

for node in range(node_count):
    # neighbor_indices[:, node] =
    print(
        np.fromiter(
            np.ndenumerate(affinity_matrix[node,:]),
            dtype=(np.intp, np.float64),
            count=node_count,
        )#.sort(
     #    reverse=True, key=lambda x: x[1]
     #    )[1 :: k + 1][0]
    )

错误情况 #

1. 使用dtype=(np.intp, np.float64)或dtype=(int, np.float64)时，抛出ValueError: setting an array element with a sequence.，原因是试图将多值序列赋值给单个数组元素。
2. 查看np.ndenumerate输出结构：

   for x in np.ndenumerate(affinity_matrix[node, :]):
       print(x)
       print(type(x), " ", type(x[0]), " ", type(x[0][0]))
   

   输出结果：

   ((990,), 0.9958856990164133)
   <class 'tuple'>   <class 'tuple'>   <class 'int'>
   

3. 尝试dtype=((int,), np.float64)时，抛出TypeError: Tuple must have size 2, but has size 1。

解决方案 #

1. 正确结合np.fromiter与np.ndenumerate #

可以结合使用，但需要匹配np.ndenumerate的输出结构设置结构化dtype：每个输出元素是((索引元组), 值)，因此dtype需要定义为包含两个字段的结构化类型，分别对应索引元组和数值。

修正后的代码示例：

import numpy as np

# 假设参数已定义：affinity_matrix(n×n), node_count=n, k=邻居数
neighbor_indices = np.zeros((node_count, k), dtype=np.intp)

for node in range(node_count):
    # 用结构化dtype匹配ndenumerate的输出结构
    enum_array = np.fromiter(
        np.ndenumerate(affinity_matrix[node, :]),
        dtype=[('index_tuple', 'i8', 1), ('value', 'f8')],
        count=node_count
    )
    # 按值降序排序（如果是距离则改为升序，去掉[::-1]）
    sorted_array = np.sort(enum_array, order='value')[::-1]
    # 提取前k个索引（从元组中取出单个元素）
    top_k_indices = sorted_array[:k]['index_tuple'][:, 0]
    neighbor_indices[node] = top_k_indices

2. 更高效的替代方法（推荐） #

实际上无需使用np.ndenumerate，直接用np.argsort可以更简洁高效地实现需求：

import numpy as np

neighbor_indices = np.zeros((node_count, k), dtype=np.intp)
for node in range(node_count):
    # 对该行的值排序，获取索引（降序取前k，若为距离则升序，去掉[::-1]）
    sorted_indices = np.argsort(affinity_matrix[node, :])[::-1][:k]
    neighbor_indices[node] = sorted_indices

这个方法避免了复杂的迭代器和结构化数组操作，性能更优。

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