用索引的numpy数组索引numpy数组 [英] Indexing a numpy array with a numpy array of indexes
问题描述
我有一个3D numpy数组data
和另一个索引数组pos
(索引本身就是一个numpy数组,这使后者成为2D数组):
I have a 3D numpy array data
and another array pos
of indexes (an index is a numpy array on its own, which makes the latter array a 2D array):
import numpy as np
data = np.arange(8).reshape(2, 2, -1)
#array([[[0, 1],
# [2, 3]],
#
# [[4, 5],
# [6, 7]]])
pos = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0]])
#array([[1, 1, 0],
# [0, 1, 0],
# [1, 0, 0]])
我想使用pos
的索引选择和/或变异data
中的元素.我可以使用for
循环或列表理解来进行选择:
I want to select and/or mutate the elements from data
using the indexes from pos
. I can do the selection using a for
loop or a list comprehension:
[data[tuple(i)] for i in pos]
#[6, 2, 4]
data[[i for i in pos.T]]
#array([6, 2, 4])
但这似乎不是一种麻木的方式.是否有针对此问题的矢量化numpy解决方案?
But this does not seem to be a numpy way. Is there a vectorized numpy solution to this problem?
推荐答案
您可以将pos
拆分为3个单独的数组和索引,就像这样—
You can split pos
into 3 separate arrays and index, like so—
>>> i, j, k = pos.T
>>> data[i, j, k]
array([6, 2, 4])
在这里,pos
中的列数与data
的深度相对应.只要您要处理3D矩阵,就很好地获得i
,j
和k
的复杂性.
Here, the number of columns in pos
correspond to the depth of data
. As long as you're dealing with 3D matrices, getting i
, j
, and k
well never get more complicated than this.
在python-3.6 +上,您可以将其缩短为-
On python-3.6+, you can shorten this to—
>>> data[[*pos.T]]
array([6, 2, 4])
这篇关于用索引的numpy数组索引numpy数组的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!