python numpy数组切片

Question

I have an 2d array, A that is 6x6. I would like to take the first 2 values (index 0,0 and 0,1) and take the average of the two and insert the average into a new array that is half the column size of A (6x3) at index 0,0. Then i would get the next two indexes at A, take average and put into the new array at 0,1.

The only way I know how to do this is using a double for loop, but for performance purposes (I will be using arrays as big as 3000x3000) I know there is a better solution out there! Thanks!

Answer 1

A very useful feature of numpy arrays is that they can be reshaped and viewed in many different ways, and by doing so, you can make certain operations very easy.

Since you want to pair every two items, it makes sense to reshape the 6x6 array into a 18x2 array:

import numpy as np

arr=np.arange(36).reshape(6,6)
print(arr)
# [[ 0  1  2  3  4  5]
#  [ 6  7  8  9 10 11]
#  [12 13 14 15 16 17]
#  [18 19 20 21 22 23]
#  [24 25 26 27 28 29]
#  [30 31 32 33 34 35]]
arr2=arr.reshape(-1,2)
print(arr2)
# [[ 0  1]
#  [ 2  3]
#  [ 4  5]
#  [ 6  7]
#  [ 8  9]
#  [10 11]
#  [12 13]
#  [14 15]
#  [16 17]
#  [18 19]
#  [20 21]
#  [22 23]
#  [24 25]
#  [26 27]
#  [28 29]
#  [30 31]
#  [32 33]
#  [34 35]]

Now taking the average is easy:

means=arr2.mean(axis=1)
print(means)
# [  0.5   2.5   4.5   6.5   8.5  10.5  12.5  14.5  16.5  18.5  20.5  22.5
#   24.5  26.5  28.5  30.5  32.5  34.5]

And finally, we just reshape the array to be 6x3:

means=means.reshape(6,-1)
print(means)
# [[  0.5   2.5   4.5]
#  [  6.5   8.5  10.5]
#  [ 12.5  14.5  16.5]
#  [ 18.5  20.5  22.5]
#  [ 24.5  26.5  28.5]
#  [ 30.5  32.5  34.5]]

or, as a 1-liner:

means=arr.reshape(-1,2).mean(axis=1).reshape(6,-1)

PS: reshaping is a very quick operation, since it is returning a view, not a copy of the original array. All that is changed is the dimensions and the strides. All that remains is one call to the mean method. Thus, this solution should be about a quick as possible using numpy.

Answer 2

I don't think there is a better solution, unless you have some extra information about what's in those arrays. If they're just random numbers, you have to do (n^2)/2 calculations, and your algorithm is reflecting that, running in O((n^2)/2).