如何根据指数给定沿给定轴的元素?

时间:2017-07-02 07:23:28

标签: python arrays numpy

我有一个3D数组,我需要在最后一个轴上“挤”它,这样我得到一个2D数组。我需要以下面的方式做到这一点。对于前两个维度的索引的每个值,我知道应从中获取值的第三维索引的值。

例如,我知道如果i1 == 2i2 == 7那么i3 == 11。这意味着out[2,7] = inp[2,7,11]。从前两个维度到第三个维度的映射在另一个2D阵列中给出。换句话说,我有一个数组,其位置2,7我有11作为值。

所以,我的问题是如何组合这两个数组(3D和2D)来获得输出数组(2D)。

4 个答案:

答案 0 :(得分:2)

In [635]: arr = np.arange(24).reshape(2,3,4)
In [636]: idx = np.array([[1,2,3],[0,1,2]])


In [637]: I,J = np.ogrid[:2,:3]
In [638]: arr[I,J,idx]
Out[638]: 
array([[ 1,  6, 11],
       [12, 17, 22]])
In [639]: arr
Out[639]: 
array([[[ 0,  1,  2,  3],   # 1
        [ 4,  5,  6,  7],   # 6
        [ 8,  9, 10, 11]],  # ll

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])

I,J一起广播以选择匹配idx的一组(2,3)值:

In [640]: I
Out[640]: 
array([[0],
       [1]])
In [641]: J
Out[641]: array([[0, 1, 2]])

这是对3d更简单的问题的概括 - 从每行中选择一个项目:

In [649]: idx
Out[649]: 
array([[1, 2, 3],
       [0, 1, 2]])
In [650]: idx[np.arange(2), [0,1]]
Out[650]: array([1, 1])

事实上,我们可以将3d问题转换为2d问题:

In [655]: arr.reshape(6,4)[np.arange(6), idx.ravel()]
Out[655]: array([ 1,  6, 11, 12, 17, 22])

概括原案例:

In [55]: arr = np.arange(24).reshape(2,3,4)                                     
In [56]: idx = np.array([[1,2,3],[0,1,2]])                                      
In [57]: IJ = np.ogrid[[slice(i) for i in idx.shape]]                           
In [58]: IJ                                                                     
Out[58]: 
[array([[0],
        [1]]), array([[0, 1, 2]])]
In [59]: (*IJ,idx)                                                              
Out[59]: 
(array([[0],
        [1]]), array([[0, 1, 2]]), array([[1, 2, 3],
        [0, 1, 2]]))
In [60]: arr[_]                                                                 
Out[60]: 
array([[ 1,  6, 11],
       [12, 17, 22]])

关键是将IJ数组列表与idx组合在一起,以创建一个新的索引元组。如果idx不是最后一个索引,那么构造元组会有点麻烦,但它仍然是可能的。 E.g。

In [61]: (*IJ[:-1],idx,IJ[-1])                                                  
Out[61]: 
(array([[0],
        [1]]), array([[1, 2, 3],
        [0, 1, 2]]), array([[0, 1, 2]]))
In [62]: arr.transpose(0,2,1)[_]                                                
Out[62]: 
array([[ 1,  6, 11],
       [12, 17, 22]])

如果更容易将arr转置到idx维度,则最后。关键是索引操作采用索引数组的元组,这些数组相互广播以选择特定项。 这就是ogrid正在做的事情,创建使用idx的数组。

答案 1 :(得分:1)

inp = np.random.random((20, 10, 5)) # simulate some input
i1, i2 = np.indices(inp.shape[:2])
i3 = np.random.randint(0, 5, size=inp.shape) # or implement whatever mapping
                                             # you want between (i1,i2) and i3
out = inp[(i1, i2, i3)]

有关详细信息,请参阅https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#integer-array-indexing

答案 2 :(得分:0)

使用numpy.einsum

这可以通过array indexingnumpy.einsum的使用:

的组合来实现
>>> numpy.einsum('ijij->ij', inp[:, :, indices])

inp[:, :, indices]创建一个四维数组,其中对于前两个索引(前两个维度)中的每个索引,索引数组的所有索引都应用于第三个维度。因为索引数组是二维的,所以得到4D。但是,您只希望索引数组的那些索引对应于前两个维度的索引。然后通过使用字符串ijij->ij来实现。这告诉einsum您只想选择第1轴和第3轴以及第2轴和第4轴的索引相似的元素。由于索引数组添加了最后两个维度(第3个和第4个维度),因此类似于仅为index[i, j]的第三个维度选择索引inp

请注意,此方法可以真正炸掉内存消耗。特别是如果inp.shape[:2]远大于inp.shape[2],那么inp[:, :, indices].size将大约为inp.size ** 2

手动构建索引

首先我们准备新的索引数组:

>>> idx = numpy.array(list(
...     numpy.ndindex(*inp.shape[:2], 1)  # Python 3 syntax
... ))

然后我们更新对应于第三轴的列:

>>> idx[:, 2] = indices[idx[:, 0], idx[:, 1]]

现在我们可以选择元素并简单地重塑结果:

>>> inp[tuple(idx.T)].reshape(*inp.shape[:2])

使用numpy.choose

注意: numpy.choose允许从中选择的轴的最大大小为32。

根据this answernumpy.choose的文档,我们还可以使用以下内容:

# First we need to bring the last axis to the front because
# `numpy.choose` chooses from the first axis.
>>> new_inp = numpy.moveaxis(inp, -1, 0)
# Now we can select the elements.
>>> numpy.choose(indices, new_inp)

虽然文档不鼓励使用单个数组作为第二个参数(选项)

  

为了减少误解的可能性,即使名义上支持以下“滥用”,选择既不应该也不应被视为单个数组,即最外层的序列容器应该是列表或元组。

这似乎只是为了防止误解:

  

选项数组序列

     

选择数组。 a 并且所有选项必须可以播放到相同的形状。如果选项本身就是一个数组(不推荐),那么它的最外层维度(即与choices.shape[0]对应的维度)将被视为定义“序列”。

因此,从我的观点来看,只要有人意识到他们正在做的事情,那么使用numpy.choose就没有错。

答案 3 :(得分:-1)

我相信应该这样做:

for i in range(n):
    for j in range(m):
        k = index_mapper[i][j]
        value = input_3d[i][j][k]
        out_2d[i][j] = value