使用另一个数组将超过某个索引的 numpy 数组值设置为 0

Question

我有两个数组，一个是浮点数，一个是整数

arr1 = np.asarray([[1.5, 0.75, 0.2],
            [0.3, 1.8, 4.2]])
arr2 = np.asarray([2, 1])

我需要修改 arr1 以便

arr1 [0, arr2[0]:] = 0
arr1 [1, arr2[1]:] = 0

获得

array([[1.5 , 0.75, 0.  ],
       [0.3 , 0.  , 0.  ]])

即

arr2

的第 n 个元素指示

arr1

的第 n 行中非零元素的数量。

如何将其矢量化为大型数组，并将其扩展到更高的维度？说

rng = np.random.default_rng()
arr1 = rng.uniform(size=(100, 10, 1000))
arr2 = rng.integers(10, size=(100, 1000))

我如何更改

arr1

（或根据实现创建一个新的

arr3

），以便每个切片

arr1[:, :, i], arr2[:, i]

遵循以上内容？

最终实现量级为 100,000 x 100 x 10,000。

Answer 1

我认为诀窍是使用

arr2

和代表索引的数字

np.arange

创建一个掩码。在你的玩具示例中：

>>> arr2[:, None] > np.arange(3)
array([[ True,  True, False],
       [ True, False, False]])

它创建了一个掩码，其中包含要保留的值以及用零替换的值。然后您可以：

arr3 = np.where(mask, arr1, 0)

，也可以直接就地修改

arr1

。

在你的最后一个例子中，我不完全确定你想沿着哪个维度用零替换切片，或者最终的顺序意味着什么。所以这就是一般的技巧：

# Let's say we want to replace slices along the 'D' axis.
arr1 = rng.uniform(size=(A, B, C, D, E, F))

# The slice indices only make sense if we choose integers that are smaller than D
arr2 = rng.integers(D, size=(A, B, C, E, F))

# Need to insert an empty axis in `arr2` with respect to the dimension in `arr1`
# along which you want to replace the slices.
# We do the opposite for the `np.arange` indices.
# That's because the broadcasting plays out nicely.
mask = arr2[:, :, :, None, :, :] > np.arange(D)[None, None, None, :, None, None]

arr3 = np.where(mask, arr1, 0)