numpy中的3d数组访问

Question

I have NxMx3 numpy array and want to get access to subarrays of size 3. For example I want instead of

arr = [[[...]]]
for i in range(N):
    for j in range(M):
        b = do_something(arr[i][j])

write

map(lambda x: do_something(x), ???) # x - is array of size 3

How can I do this?

The function do_something takes arrays of size 3 and returns a scalar, and I want to get the array of the results by applying the function to each length-3 subarray of my input.

Answer 1

If possible, you should manually vectorize your original function, because making use of vectorized arithmetic would be the most efficient solution.

If you don't want to or can't do that, you can use numpy.vectorize to use your function that works along a single dimension and generalize it to higher-dimensional arrays. it probably won't be faster than looping manually ( vectorize is essentially a wrapper for a for loop), but at least it will give you a more straightforward interface to calling your function.

Example:

import numpy as np

def foo(x):
    ''''sum three numbers from input'''
    return x[0]+x[1]+x[2] # otherwise x.sum(), of course

foo_vector = np.vectorize(foo,signature='(n)->()')

# try with dummy input
arr = np.random.rand(2,4,3)
#print(foo(arr)) # leads to an error
print(foo_vector(arr).shape)
# (2, 4)
print(np.allclose(foo_vector(arr),arr.sum(axis=-1)))
# True

As you can see from the above, vectorize handles additional leading dimensions of your array, and applies your function appropriately. Since you exactly want to apply your function along the last dimension, this should work like a charm, and applying foo_vector to an array of shape (M,N,3) should return an array of shape (M,N) with the result.