如何基于条件和列索引对NumPy二维数组元素执行指定计算

你可以利用NumPy的广播机制直接构造列索引数组，结合np.where完成需求，完整实现如下：

import numpy as np

# 输入数组
ar = np.array([[  0,  -5,  -5, -11,  -9],
               [  5,   0,   0,  -6,  -4],
               [ 10,   5,   5,  -1,   1]])

# 生成对应列的索引数组
col_idx = np.arange(ar.shape[1])
# 条件计算：元素为0时返回10的对应列索引次方，否则返回0
res = np.where(ar == 0, 10 ** col_idx, 0)

运行后输出的数值结果为：

array([[  1,   0,   0,   0,   0],
       [  0,  10, 100,   0,   0],
       [  0,   0,   0,   0,   0]])

如果你需要保留10^i的字符串展示形式，和你给出的期望输出格式完全对齐，可以修改为字符串计算逻辑：

res_str = np.where(ar == 0, [f'10^{i}' for i in col_idx], '0')

输出结果为：

array([['10^0', '0', '0', '0', '0'],
       ['0', '10^1', '10^2', '0', '0'],
       ['0', '0', '0', '0', '0']], dtype='<U4')

逻辑说明 #

长度等于数组列数的一维列索引数组在和二维布尔矩阵运算时，NumPy会自动将列索引沿行方向广播，每个元素都能匹配到对应列的索引值，不需要手动循环或者扩展数组维度，计算效率远高于遍历数组的实现方案。

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