如何用numpy.argpartition结果构建掩码以保留距离矩阵每行k个最小值？

1. 正确创建掩码并生成密集矩阵 #

你的问题出在np.isin(data, nn_indices)的用法上——data是距离矩阵的数值，nn_indices是列索引，两者类型和含义完全不同，自然得不到正确的掩码。

正确的做法是利用行索引和得到的列索引来标记需要保留的位置：

import numpy as np

# 示例输入数据
data = np.array([[4, 6, 1, 3],
                 [1, 5, 6, 7],
                 [4, 7, 2, 3],
                 [7, 1, 8, 2]])
n = data.shape[0]
k = 2

# 获取每行前k小值的列索引
nn_indices = np.argpartition(data, k - 1)[:, :k]

# 创建全False的掩码
mask = np.zeros_like(data, dtype=bool)
# 为每行的指定列索引位置设为True
mask[np.arange(n)[:, None], nn_indices] = True

# 生成保留k个最小值的矩阵（其余为0）
result = np.where(mask, data, 0)
print(result)

运行后会得到你期望的输出：

[[0 0 1 3]
 [1 5 0 0]
 [0 0 2 3]
 [0 1 0 2]]

2. 直接构建Scipy稀疏矩阵（csr_matrix） #

如果要直接生成稀疏矩阵，不需要先构建密集矩阵，可以提取保留位置的行坐标、列坐标、对应数值，然后传入csr_matrix：

from scipy.sparse import csr_matrix

# 提取行索引：每行重复k次（因为每行保留k个值）
rows = np.repeat(np.arange(n), k)
# 提取列索引：把nn_indices展平成一维数组
cols = nn_indices.flatten()
# 提取对应位置的数值
values = data[rows, cols]

# 构建csr稀疏矩阵
sparse_mat = csr_matrix((values, (rows, cols)), shape=(n, n))

# 可以转为密集矩阵验证结果
print(sparse_mat.toarray())

这样得到的稀疏矩阵和上面的密集矩阵完全对应，同时节省内存（适合大矩阵场景）。

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