当仅给出范围的索引时，索引数组列的范围

Question

我正在寻找一种有效的方法来索引具有多个范围的 numpy 数组的列，当只给出所需范围的索引时。

例如，给定以下数组和范围大小 r_size=3:

import numpy as np
arr = np.arange(18).reshape((2,9))

array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8],
       [ 9, 10, 11, 12, 13, 14, 15, 16, 17]])

这意味着总共有 3 组范围 [r0, r1, r2]，其数组中的元素分布为：

[[r0_00, r0_01, r0_02, r1_00, r1_01, r1_02, r2_00, r2_01, r2_02]
 [r0_10, r0_11, r0_12, r1_10, r1_11, r1_12, r2_10, r2_11, r2_12]]

因此，如果我想访问范围 r0 和 r2，那么我将获得：

arr    = np.arange(18).reshape((2,9))
r_size = 3
ranges = [0, 2]
# --------------------------------------------------------
# Line that index arr, with the variable ranges... Output:
# --------------------------------------------------------
array([[ 0,  1,  2,  6,  7,  8],
       [ 9, 10, 11, 15, 16, 17]])

我发现的最快方法如下：

import numpy as np
from itertools import chain

arr    = np.arange(18).reshape((2,9))
r_size = 3
ranges = [0,2]

arr[:, list(chain(*[range(r_size*x,r_size*x+r_size) for x in ranges]))]

array([[ 0,  1,  2,  6,  7,  8],
       [ 9, 10, 11, 15, 16, 17]])

但是我不确定是否可以在速度方面进行改进。

提前致谢！

Answer 1

您可以先将数组分成 r_size 个块：

>>> splits = np.split(arr, r_size, axis=1)
[array([[ 0,  1,  2],
        [ 9, 10, 11]]), 
 array([[ 3,  4,  5],
        [12, 13, 14]]), 
 array([[ 6,  7,  8],
        [15, 16, 17]])]

堆叠 np.stack 和 select 正确的 ranges:

>>> stack = np.stack(splits)[ranges]
array([[[ 0,  1,  2],
        [ 9, 10, 11]],

       [[ 6,  7,  8],
        [15, 16, 17]]])

并在 axis=1 上与 np.hstack or np.concantenate 水平连接：

>>> np.stack(stack)
array([[ 0,  1,  2,  6,  7,  8],
       [ 9, 10, 11, 15, 16, 17]])

整体看起来像：

>>> np.hstack(np.stack(np.split(arr, r_size, axis=1))[ranges])
array([[ 0,  1,  2,  6,  7,  8],
       [ 9, 10, 11, 15, 16, 17]])

或者，您可以专门使用 np.reshapes，这样会更快：

初始整形：

>>> arr.reshape(len(arr), -1, r_size)
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

索引 ranges:

>>> arr.reshape(len(arr), -1, r_size)[:, ranges]
array([[[ 0,  1,  2],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [15, 16, 17]]])

然后，重塑回最终形态：

>>> arr.reshape(len(arr),  -1, r_size)[:, ranges].reshape(len(arr), -1)

Answer 2

您将不可避免地需要复制数据以在连续数组中获得所需的结果。尽管为了提高效率，我建议尽量减少复制数据的次数。任何类型的整形操作都可以用 np.lib.stride_tricks.as_strided.

表示

假设原数组包含64位整数，那么每个元素都是8个字节，按某种形状排列：

import numpy as np
arr = np.arange(18).reshape((2,9))
arr.shape, arr.strides

输出：

((2, 9), (72, 8))

所以每列跳过 8 个字节，每行跳过 72 个字节。 arr.reshape(len(arr), -1, r_size)可以表示为：

np.lib.stride_tricks.as_strided(arr, (2,3,3), (72,24,8))

输出：

array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

而arr.reshape(len(arr), -1, r_size)[:, ranges]可以表示为：

np.lib.stride_tricks.as_strided(arr, (2,2,3), (72,24*2,8))

输出：

array([[[ 0,  1,  2],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [15, 16, 17]]])

到目前为止，我们只更改了数组的元数据，这意味着没有复制任何数据。此操作的性能成本接近于零。但是要获得最终数组，您需要以某种方式复制数据：

np.lib.stride_tricks.as_strided(arr, (2,2,3), (72,24*2,8)).reshape(len(arr), -1)

输出：

array([[ 0,  1,  2,  6,  7,  8],
       [ 9, 10, 11, 15, 16, 17]])

这不是一个通用的解决方案，但它可能会给您一些关于如何优化的想法。

不幸的是，我的计时并不支持这些说法，但它仍然很直观，值得对一些更大的阵列进行测试。

当仅给出范围的索引时，索引数组列的范围

Indexing ranges of columns of array when only the indexes of the ranges are given

python

indexing

numpy

range