如何使用Numpy以指定数组为参考查找矩阵中元素的索引？

给定NumPy矩阵和数组如下：

import numpy as np

A = np.array([[0.        , 0.        , 0.        , ..., 0.        , 0.        ,
    0.        ],
   [0.        , 0.28867513, 0.28867513, ..., 0.        , 0.        ,
    0.        ],
   [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
    0.        ],
   [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
    0.        ],
   [0.        , 0.13363062, 0.13363062, ..., 0.        , 0.        ,
    0.        ]])

B = np.array([0.70710678, 0.66666667, 0.5       , 0.75      , 1.        ])

需要实现：

• 查找数组B中每个元素在矩阵A对应行中的索引位置，输出包含这些位置的结果
• 仅使用NumPy内置命令/逻辑，禁止用列表推导或for循环

补充说明：之前试过isin方法，但无法处理同一行存在多个相同值的情况，因此更新了问题示例。

方法1：获取所有匹配的索引对 #

利用广播机制实现逐元素匹配，再提取索引：

# 生成布尔矩阵，标记A中每个位置是否等于B对应行的元素
matches = np.equal(A, B[:, np.newaxis])
# 提取所有匹配位置的[行索引, 列索引]
all_indices = np.argwhere(matches)

all_indices会返回所有满足条件的索引对，比如如果A的第1行第1、2列元素等于B[1]，就会输出[[1,1], [1,2]]这类结果。

方法2：按B元素分组获取列索引 #

如果需要按B中每个元素分组，得到对应行的匹配列索引列表：

matches = np.equal(A, B[:, np.newaxis])
# 计算每行的匹配数量，用于拆分结果
match_counts = np.sum(matches, axis=1)
# 拆分列索引，得到每个B元素对应的匹配位置
grouped_indices = np.split(np.argwhere(matches)[:, 1], np.cumsum(match_counts)[:-1])

grouped_indices[i]就是B[i]在A第i行中的所有列索引数组。

方法3：生成同形状的位置标记矩阵 #

如果需要输出和A同形状的矩阵，标记每个位置是否为对应B元素的匹配索引（无匹配用-1标记）：

matches = np.equal(A, B[:, np.newaxis])
# 初始化结果矩阵，默认值为-1
result_matrix = np.full_like(A, -1, dtype=int)
# 获取匹配位置的行、列索引
row_idx, col_idx = np.where(matches)
# 填充匹配位置的列索引
result_matrix[row_idx, col_idx] = col_idx

比如A第1行第1、2列匹配B[1]，则result_matrix[1,1]和result_matrix[1,2]会被设为1和2，其余位置为-1。

关键说明 #

• B[:, np.newaxis]将一维数组B转为列向量，通过广播实现和A的逐行逐元素比较，完美适配“B第i个元素对应A第i行”的匹配逻辑
• 上述方法均未使用循环或列表推导，完全基于NumPy内置函数实现，同时能处理同一行多个相同值的场景

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