Python numpy 2D array sum over certain indices

问题

There is a 2-d array like this:

img = [
  [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
  [[2, 2, 2], [3, 2, 3], [6, 7, 6]],
  [[9, 8, 1], [9, 8, 3], [9, 8, 5]]
]

And i just want to get the sum of certain indices which are like this:

indices = [[0, 0], [0, 1]] # which means img[0][0] and img[0][1]
# means here is represents

There was a similar ask about 1-d array in stackoverflow in this link, but it got a error when I tried to use print(img[indices]). Because I want to make it clear that the element of img are those which indicates by indices, and then get the mean sum of it.

Expected output

[5, 7, 9]

回答1:

Use NumPy:

import numpy as np

img = np.array(img)
img[tuple(indices)].sum(axis = 0)
#array([5, 7, 9])

回答2:

If the result would be [5, 7, 9] which is sum over the column of the list. Then easy:

img = np.asarray(img)
indices = [[0, 0], [0, 1]]
img[(indices)].sum(axis = 0)

Result:

array([5, 7, 9])

回答3:

When you supply a fancy index, each element of the index tuple represents a different axis. The shape of the index arrays broadcasts to the shape of the output you get.

In your case, the rows of indices.T are the indices in each axis. You can convert them into an index tuple and append slice(None), which is the programmatic equivalent of :. You can take the mean of the resulting 2D array directly:

img[tuple(indices.T) + (slice(None),)].sum(0)

Another way is to use the splat operator:

img[(*indices.T, slice(None))].sum(0)

来源：https://stackoverflow.com/questions/65768227/python-numpy-2d-array-sum-over-certain-indices