如何从numpy.where()提取单个位置值并添加至数组，求解匹配索引

嘿，我来帮你搞定这个Numpy数组的索引问题，顺便把np.where()的用法给你讲明白～

首先先澄清个小细节：你用np.max(a, axis=0)生成的b其实应该是一维数组（形状是(列数,)），不是你写的二维嵌套列表[[...]]——如果确实是二维的话，先做b = b.flatten()或者直接用b = np.max(a, axis=0)拿到一维数组，后续操作会更顺畅。

一、换掉嵌套循环，高效找到匹配列最大值的元素索引 #

你现在用嵌套循环的方式在数组大的时候效率会很低，Numpy的向量化操作能一步搞定，而且速度快很多：

情况1：每列可能有多个最大值（需要找到所有匹配的索引） #

我们可以先利用Numpy的广播机制，让b自动扩展成和a同形状的数组，然后逐元素比较，再用np.where()拿到所有符合条件的索引：

import numpy as np

# 搞个示例数组方便演示
a = np.array([[1.0, 0.5, 0.3],
              [0.0, 0.4, 0.6],
              [1.0, 0.5, 0.0]])  # 第0列和第1列都有两个最大值
b = np.max(a, axis=0)  # 结果是array([1. , 0.5, 0.6])

# 生成布尔数组，标记哪些元素等于对应列的最大值
matches = a == b
# np.where返回两个数组：行索引数组、列索引数组
row_indices, col_indices = np.where(matches)

# 如果想按列整理索引（比如每列对应的所有行索引），可以这么做：
locn = []
for col in range(a.shape[1]):
    # 筛选出当前列的所有匹配行索引
    cols_matches = row_indices[col_indices == col]
    locn.append(cols_matches.tolist())

print(locn)
# 输出：[[0, 2], [0, 2], [1]]，对应每列所有最大值的行位置

情况2：每列最大值唯一（只需要第一个出现的索引） #

这种情况更简单，直接用np.argmax()就能拿到每列最大值的行索引，一步到位：

# 用刚才的示例数组（如果最大值唯一的话）
row_indices = np.argmax(a, axis=0)
# 结果是array([0, 0, 1])，对应每列最大值的行位置

# 要是想转换成(行, 列)的索引对列表：
locn = list(zip(row_indices, np.arange(a.shape[1])))
print(locn)
# 输出：[(0, 0), (0, 1), (1, 2)]

二、从np.where()获取单个位置并添加到数组 #

np.where()返回的是两个一维数组，分别对应所有匹配位置的行和列索引。如果你想取出单个位置，直接索引这两个数组就行：

比如拿第一个匹配的位置：

row_indices, col_indices = np.where(a == b)
# 取出第一个匹配的行和列
first_row = row_indices[0]
first_col = col_indices[0]

# 把这个位置添加到另一个numpy数组里：
# 先初始化一个空数组
index_array = np.empty((0, 2), dtype=int)
# 添加第一个位置
index_array = np.append(index_array, [[first_row, first_col]], axis=0)
# 再添加第二个位置
index_array = np.append(index_array, [[row_indices[1], col_indices[1]]], axis=0)

print(index_array)
# 输出：[[0 0]
#        [2 0]]

如果想逐个遍历所有匹配位置并添加，用zip配对行和列索引就行：

result_list = []
for r, c in zip(row_indices, col_indices):
    # 把每个位置作为元组加到列表里
    result_list.append((r, c))

# 转换成numpy数组的话：
result_np = np.array(result_list)

这样操作就比嵌套循环高效多了，而且符合Numpy的向量化编程思路～

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