从大型2D NumPy数组中提取大小可变的子行

时间:2018-08-10 02:27:24

标签: python numpy submatrix

我有一个大小为3*10的NumPy数组,我想从每一行中提取大小不同的子行。子行以变化的像素大小在中间像素居中。然后,我取每个子行的平均数。我在下面有一个伪示例:

import numpy as np
arr = np.arange(1,31).reshape((3,10))
pixel_size = np.array([2,3,1])
## the subrow centers in the middle of the array, index 5
mask = [[5-2:5+2],[5-3:5+3],[5-1:5+1]] ## index for each row
### submatrix = arr[;,mask]
submatrix = [[3,4,5,6],[12,13,14,15,16,17],[24,25]]
## output = np.mean(submatrix, axis=1) output is the average number of each row in the submatrix
output = [4.5,14.5,24.5]

如果我有超过1000万行,那么如何快速处理这种情况。

1 个答案:

答案 0 :(得分:3)

您可以使用列表推导和索引切片来实现:

import numpy as np

arr = np.arange(1,31).reshape((3,10))
pixel_size = np.array([2,3,1])

middle_ind = int(arr.shape[1]/2.)
print middle_ind
sub_arr = [arr[i,middle_ind - pixel_size[i]:middle_ind + pixel_size[i]] for i in range(len(pixel_size))]
print('sub_arr: ', sub_arr)
output = [np.mean(item) for item in sub_arr]
print('output: ', output)

> sub_arr: [array([4, 5, 6, 7]), array([13, 14, 15, 16, 17, 18]), array([25, 26])]
> output: [5.5, 15.5, 25.5]

您的子矩阵是一个列表,而不是数组,因此向量化操作更加困难。您可能需要考虑重组代码以利用矩阵运算。