在使用Pythonic向量化代码时，如何在2D numpy数组中索引不同大小的范围？

Question

在使用Pythonic向量化代码时，如何在2D numpy数组中索引不同大小的范围？

8

我有一个numpy的2D数组，我想根据列索引选择不同大小的数组范围。这是输入数组的示例。

[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]
 [12 13 14]]

那么，列表 b = [4,3,1] 确定了每个列切片的不同区间大小，以便我们得到以下数组

[0 3 6 9]
[1 4 7]
[2]

我们可以将它们连接和展平以获得最终所需的输出。

[0 3 6 9 1 4 7 2]

目前，为了完成这个任务，我正在使用以下代码

slices = []
for i in range(a.shape[1]):
    slices.append(a[:b[i],i])

c = np.concatenate(slices)

如果可能的话，我希望你能将其转换为Python风格的格式。

奖励：相同的问题，但现在考虑到b决定行切片而不是列。

- xicocaio

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Divakar · Accepted Answer

我们可以使用“广播”来生成一个合适的掩码，然后“屏蔽”会完成这项工作。

In [150]: a
Out[150]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

In [151]: b
Out[151]: [4, 3, 1]

In [152]: mask = np.arange(len(a))[:,None] < b

In [153]: a.T[mask.T]
Out[153]: array([0, 3, 6, 9, 1, 4, 7, 2])

另一种掩盖的方法是 -

In [156]: a.T[np.greater.outer(b, np.arange(len(a)))]
Out[156]: array([0, 3, 6, 9, 1, 4, 7, 2])

奖金：按行切片

如果我们需要根据块大小按行进行切片，我们需要修改一些内容：

In [51]: a
Out[51]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

# slice lengths per row
In [52]: b
Out[52]: [4, 3, 1]

# Usual loop based solution :
In [53]: np.concatenate([a[i,:b_i] for i,b_i in enumerate(b)])
Out[53]: array([ 0,  1,  2,  3,  5,  6,  7, 10])

# Vectorized mask based solution :
In [54]: a[np.greater.outer(b, np.arange(a.shape[1]))]
Out[54]: array([ 0,  1,  2,  3,  5,  6,  7, 10])