输入数组相等时索引矩阵元素

时间:2019-01-29 06:30:07

标签: python arrays numpy matrix indexing

我有一个矩阵,我希望能够在不使用循环的情况下用两个数组索引某些元素时更改它们的值

例如

import numpy as np

A = np.array([[0, 0, 0, 0],
              [0, 0, 0, 0],
              [0, 0, 0, 0]])
b = np.array([0, 1, 2])
c = np.array([2, 3, 1, 1])

如果b中的索引i等于c中的索引j,则访问A [i,j]

我尝试过

A[b, c] = 1

A[b == c] = 1

但是我没有得到正确的结果 我相信结果应该是

A = array([[0, 0, 0, 0],
           [0, 0, 1, 1],
           [1, 0, 0, 0]])

1 个答案:

答案 0 :(得分:1)

您可以使用numpy.meshgrid将1D索引数组扩展为2D索引数组:

import numpy as np

A = np.array([[0, 0, 0, 0],
              [0, 0, 0, 0],
              [0, 0, 0, 0]])
b = np.array([0, 1, 2])
c = np.array([2, 3, 1, 1])

C,B = np.meshgrid(c,b)
A[B==C] = 1

print(A)

给予

[[0 0 0 0]
 [0 0 1 1]
 [1 0 0 0]]