切片并用1d列索引数组填充2d数组

Question

首先，这是我想要做的一维模拟。

假设我有一个0s的1d数组，并且我想用1替换索引2以后的每个0。我可以这样做：

import numpy as np

x = np.array([0,0,0,0])
i = 2
x[i:] = 1
print(x)  # [0 0 1 1]

现在，我正在尝试找出此操作的2d版本。首先，我有一个5x4的0数组，例如

foo = np.zeros(shape = (5,4))

[[0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]]

和相应的列索引的5个元素数组，例如

fill_locs = np.array([0, 3, 1, 1, 2])

对于foo的每一行，我想用1s填充i:列，其中i是fill_locs给定的索引。 foo[fill_locs.reshape(-1, 1):] = 1感觉不错，但是不起作用。

我想要的输出应该像

expected_result = np.array([
    [1, 1, 1, 1],
    [0, 0, 0, 1],
    [0, 1, 1, 1],
    [0, 1, 1, 1],
    [0, 0, 1, 1],
])

Answer 1

您不需要切片，也不需要创建原始数组。您可以通过广播比较来完成所有这些操作。

---

a = np.array([0, 3, 1, 1, 2])
n = 4

(a[:, None] <= np.arange(n)).view('i1')

array([[1, 1, 1, 1],
       [0, 0, 0, 1],
       [0, 1, 1, 1],
       [0, 1, 1, 1],
       [0, 0, 1, 1]], dtype=int8)

---

或使用less_equal.outer

np.less_equal.outer(a, np.arange(n)).view('i1')