将多行切片设置为值,每行一个唯一的切片



给定一个2d数组,我可以将行切片设置为特定值

import numpy as np
a = np.zeros(25).reshape(5,-1).astype(int)
a[0][2:4] = 1.0
a
array([[0, 0, 1, 1, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]])

我试图将多个行切片设置为一个特定的值,每行都有一个唯一的切片。

我有两个阵列中切片的开始和结束索引

starts = np.array([2, 0, 1, 3, 2])
ends = np.array([5, 3, 4, 5, 4])

但我似乎想不出一种方法来将2d数组的这些切片设置为特定值

a[starts:ends] = 1

TypeError: only integer scalar arrays can be converted to a scalar index中的结果

如果目标数组的最后一个维度很大,那么使用基本的Python循环相对高效,因为与填充数组相比,Python循环的开销较小。否则,AFAIK-Numpy无法提供任何有效执行此操作的方法(主要是因为切片的大小可变(。下面是一个带有Python循环的基本代码:

for i, start, end in zip(range(starts.size), starts.tolist(), ends.tolist()):
a[i, start:end] = 1

如果你想要一个更快的代码,那么你可以使用Numba使循环更快。请注意,在这种情况下,您不需要调用tolist(其目的是通过不使用Numpy整数类型而使用CPython整数来使代码更快(。

Numpy有一个函数,允许您使用一个函数沿特定轴单独对数组应用操作。因此,在我的情况下,我可以将操作唯一地应用于每一行。

除了数组本身,apply_along_axis不允许将参数传递给函数,所以我首先将起始索引和结束索引连接到我的零数组,然后将它们从结果中分割出来。

import numpy as np
a = np.zeros(25).reshape(5,-1).astype(int)
starts = np.array([2, 0, 1, 3, 2])
ends = np.array([5, 3, 4, 5, 4])
startsT = np.expand_dims(starts, axis=0).transpose()
endsT = np.expand_dims(ends, axis=0).transpose()
aa = np.concatenate((a, startsT, endsT), axis=1)
def set_1s_by_slice(x):
x[x[-2]:x[-1]] = 1
return x
pen = np.apply_along_axis(set_1s_by_slice, 1, aa)
ult = pen[:,0:5]
ult
array([[0, 0, 1, 1, 1],
[1, 1, 1, 0, 0],
[0, 1, 1, 1, 0],
[0, 0, 0, 1, 1],
[0, 0, 1, 1, 0]])

从源代码来看,这可能不会比遍历行更快

https://github.com/numpy/numpy/blob/v1.22.0/numpy/lib/shape_base.py#L267-L414

似乎有列表的转换,尽管我不确定。

这似乎比使用apply_along_axis的其他答案在计算上更高效

indices = np.arange(a.shape[1])
mask = (indices >= starts[:, np.newaxis]) & (indices < ends[:, np.newaxis])
a[mask] = 1

最新更新