分配给二维 NumPy 数组的切片

Question

我想分配0给二维数组的不同长度切片。

例子：

import numpy as np    

arr = np.array([[1,2,3,4],
                [1,2,3,4],
                [1,2,3,4],
                [1,2,3,4]])

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

鉴于上述数组arr和索引idxs，您如何分配不同长度的切片。结果是：

arr = np.array([[0,2,3,4],
                [0,0,3,4],
                [0,0,0,4],
                [0,2,3,4]])

这些不起作用

slices = np.array([np.arange(i) for i in idxs])
arr[slices] = 0

---

arr[:, :idxs] = 0

Answer 1

您可以使用广播比较来生成掩码，并相应地索引到arr：

arr[np.arange(arr.shape[1]) <= idxs[:, None]] = 0

print(arr) 
array([[0, 2, 3, 4],
       [0, 0, 3, 4],
       [0, 0, 0, 4],
       [0, 2, 3, 4]])