python - 根据条件和列索引位置应用函数

Question

我有一个 2d numpy 数组，我想根据条件并考虑列索引位置对元素应用计算。

条件是:元素值== 0。

如果为 true，我想在该位置计算 10^i(其中 i 是轴 1 的索引)，否则为 0。

因此示例数组可以是:

ar = np.array([[  0,  -5,  -5, -11,  -9],
               [  5,   0,   0,  -6,  -4],
               [ 10,   5,   5,  -1,   1]])

结果数组应如下所示(计算指数):

[[  10^0,  0,  0, 0,  0],
 [  0,   10^1,   10^2,  0,  0],
 [ 0,   0,   0,  0,   0]]

我已经做到了:

np.where(ar==0)

鉴于条件为真，我不确定如何将列索引应用为指数。

谢谢!

Answer 1

创建测试数组:

a = np.random.choice(range(-4, 4), size = (3, 4))

# a = array([[ 0, -1, -1, -2],
#            [ 0, -3,  0, -1],
#            [ 0,  3, -2,  0]])

查找 a 为零的索引:

mask = a == 0
indx = np.nonzero(mask)

# indx = (array([0, 1, 1, 2, 2]), array([0, 0, 2, 0, 3]))

更改元素:

a[mask] = 10**(indx[1])
a[~mask] = 0

# a = array([[   1,    0,    0,    0],
#            [   1,    0,  100,    0],
#            [   1,    0,    0, 1000]])

indx 是一个元组，其第二个元素存储二维数组的列索引。