使用 NumPy 在索引处设置值的一种更简洁的方法

Question

我有一个初始为零的 numpy 数组，如下所示：

v = np.zeros((5, 5))
v

array([[ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.]])

我还有一组数组idx1 和idx2。

idx1

array([[0, 3],
       [0, 4],
       [1, 3],
       [2, 4]])

idx2

array([[0, 1],
       [0, 2],
       [0, 4],
       [1, 3]])

将每对值视为行和列索引。因此，例如，在idx1 中，第一对(0, 3) 将成为v[0, 3] 的索引器，依此类推。

我想先将idx1 指定的索引处的值设置为1，然后将idx2 指定的所有索引设置为0。

另外，请注意，如果某个数组中有一对(i, j)，我想同时设置v[i, j] 和v[j, i]。

我的最终结果变成：

array([[ 0.,  0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.],
       [ 1.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  0.,  0.]])

我目前通过以下方式实现：

def set_vals(x, i, j, v):
    x[i, j] = x.T[i, j] = v

v = np.zeros((5, 5))

i1, j1 = idx1[:, 0], idx1[:, 1]
i2, j2 = idx2[:, 0], idx2[:, 1]

set_vals(v, i1, j1, 1)
set_vals(v, i2, j2, 0)

v     #  the result

但是，我相信可能有更好的方法。希望听到任何改进的想法/建议。谢谢！

Answer 1

为了寻找更“紧凑”的表达方式，我得到了这个 -

v = np.zeros((5, 5))
v[tuple(np.r_[idx1,idx1[:,::-1]].T)] = 1
v[tuple(np.r_[idx2,idx2[:,::-1]].T)] = 0

---

在 python3.6+ 上，您可以使用 * 解包运算符进一步减少此问题：

v[[*np.r_[idx1,idx1[:,::-1]].T]] = 1
v[[*np.r_[idx2,idx2[:,::-1]].T]] = 0
v

array([[ 0.,  0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.],
       [ 1.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  0.,  0.]])