当只给出所需范围的索引时，我正在寻找一种有效的方法来索引具有多个范围的 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]])
但我不确定它是否可以在速度方面有所提高。
提前致谢！
您可以首先将数组分成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并选择正确的ranges:
>>> stack = np.stack(splits)[ranges]
array([[[ 0, 1, 2],
[ 9, 10, 11]],
[[ 6, 7, 8],
[15, 16, 17]]])
并使用np.hstack或np.concantenateon水平连接axis=1：
>>> 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.reshape专门使用s ，这会更快：
初步重塑：
>>> 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)
本文收集自互联网，转载请注明来源。
如有侵权，请联系 [email protected] 删除。
我来说两句