时间戳

首页 > 解决方案 > 索引 numpy 数组，每列有两个数组，开始和结束

python - 索引 numpy 数组，每列有两个数组，开始和结束

问题描述

我有一个尺寸为 NxM 的 2D numpy 数组 A 我有两个以尺寸 Nx1 开始和结束的数组

对于 A[k,:] 给出的每一行 k，我想将 start[k] 到 end[k] 的列设置为 0：

A[:,start:end]=0

但是，我得到了错误

TypeError: only integer scalar arrays can be converted to a scalar index

如果我只是索引

A[:,start]=0

有用。但这不是我想要达到的。

每个 k 的 start[k]:end[k] 的长度不同

标签： pythonnumpyindexing

解决方案

随着开始和结束索引数组的长度 -starts以及ends较大的长度，我们可以创建这些位置的掩码 -

r = np.arange(A.shape[1])
mask = (r >= np.asarray(starts)[:,None]) & (r < np.asarray(ends)[:,None])
A[mask] = 0

样品运行 -

In [47]: np.random.seed(0)

In [48]: A = np.random.randint(10,100,(4,10))

In [49]: A
Out[49]: 
array([[54, 57, 74, 77, 77, 19, 93, 31, 46, 97],
       [80, 98, 98, 22, 68, 75, 49, 97, 56, 98],
       [91, 47, 35, 87, 82, 19, 30, 90, 79, 89],
       [57, 74, 92, 98, 59, 39, 29, 29, 24, 49]])

In [50]: starts = np.array([2,4,1,0,])

In [51]: ends = np.array([4,7,5,2])

In [52]: r = np.arange(A.shape[1])
    ...: mask = (r >= np.asarray(starts)[:,None]) & (r < np.asarray(ends)[:,None])
    ...: A[mask] = 0

In [53]: A
Out[53]: 
array([[54, 57,  0,  0, 77, 19, 93, 31, 46, 97],
       [80, 98, 98, 22,  0,  0,  0, 97, 56, 98],
       [91,  0,  0,  0,  0, 19, 30, 90, 79, 89],
       [ 0,  0, 92, 98, 59, 39, 29, 29, 24, 49]])

对于一个小的开始、结束数组或者如果A是一个非常大的数组，简单的循环可能会更好。

推荐阅读