用另一个索引多维数组的多个维度 - NumPy/Python答案

【问题标题】：Index multiple dimensions of a multi-dimensional array with another - NumPy/ Python用另一个索引多维数组的多个维度 - NumPy/Python
【发布时间】：2016-04-05 16:46:59
【问题描述】：

假设我有以下形式的张量：

import numpy as np
a = np.array([ [[1,2],
                [3,4]],
               [[5,6],
                [7,3]] 
             ])
# a.shape : (2,2,2) is a tensor containing 2x2 matrices
indices = np.argmax(a, axis=2)
#print indices
for mat in a:
    max_i = np.argmax(mat,axis=1)
    # Not really working I would like to
    # change 4 in the first matrix to -1
    #  and 3 in the last to -1
    mat[max_i] = -1

print a

现在我想做的是使用索引作为 a 上的掩码，用 -1 替换每个最大元素。有没有办法做到这一点？到目前为止，我所知道的只是使用 for 循环。

【问题讨论】：

您的括号不匹配。可以酌情更新吗？
另外，你能添加你的循环代码来提供预期的输出吗？

标签： python numpy multidimensional-array vectorization axes

【解决方案1】：

这是在3D 中使用linear indexing 的一种方法-

m,n,r = a.shape
offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
np.put(a,indices + offset,-1)

示例运行 -

In [92]: a
Out[92]: 
array([[[28, 59, 26, 70],
        [57, 28, 71, 49],
        [33,  6, 10, 90]],

       [[24, 16, 83, 67],
        [96, 16, 72, 56],
        [74,  4, 71, 81]]])

In [93]: indices = np.argmax(a, axis=2)

In [94]: m,n,r = a.shape
    ...: offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
    ...: np.put(a,indices + offset,-1)
    ...: 

In [95]: a
Out[95]: 
array([[[28, 59, 26, -1],
        [57, 28, -1, 49],
        [33,  6, 10, -1]],

       [[24, 16, -1, 67],
        [-1, 16, 72, 56],
        [74,  4, 71, -1]]])

这是linear indexing 的另一种方式，但在2D -

m,n,r = a.shape
a.reshape(-1,r)[np.arange(m*n),indices.ravel()] = -1

运行时测试和验证输出 -

In [156]: def vectorized_app1(a,indices): # 3D linear indexing
     ...:   m,n,r = a.shape
     ...:   offset = n*r*np.arange(m)[:,None] + r*np.arange(n)
     ...:   np.put(a,indices + offset,-1)
     ...: 
     ...: def vectorized_app2(a,indices): # 2D linear indexing
     ...:   m,n,r = a.shape
     ...:   a.reshape(-1,r)[np.arange(m*n),indices.ravel()] = -1
     ...:  

In [157]: # Generate random 3D array and the corresponding indices array
     ...: a = np.random.randint(0,99,(100,100,100))
     ...: indices = np.argmax(a, axis=2)
     ...: 
     ...: # Make copies for feeding into functions
     ...: ac1 = a.copy()
     ...: ac2 = a.copy()
     ...: 

In [158]: vectorized_app1(ac1,indices)

In [159]: vectorized_app2(ac2,indices)

In [160]: np.allclose(ac1,ac2)
Out[160]: True

In [161]: # Make copies for feeding into functions
     ...: ac1 = a.copy()
     ...: ac2 = a.copy()
     ...: 

In [162]: %timeit vectorized_app1(ac1,indices)
1000 loops, best of 3: 311 µs per loop

In [163]: %timeit vectorized_app2(ac2,indices)
10000 loops, best of 3: 145 µs per loop

【讨论】：

@FranciscoVargas 不应该！因此，您可以将第二种方法扩展到 4D 案例，就像这样 - p,m,n,r = a.shape; a.reshape(-1,r)[np.arange(p*m*n),indices.ravel()] = -1。

【解决方案2】：

您可以使用indices 对a 的最后一个维度进行索引，前提是您还为前两个维度指定了索引数组：

import numpy as np

a = np.array([[[1, 2],
               [3, 4]],

              [[5, 6],
               [7, 3]]])

indices = np.argmax(a, axis=2)

print(repr(a[range(a.shape[0]), range(a.shape[1]), indices]))
# array([[2, 3],
#        [2, 7]])

a[range(a.shape[0]), range(a.shape[1]), indices] = -1

print(repr(a))
# array([[[ 1, -1],
#         [ 3,  4]],

#        [[ 5,  6],
#         [-1, -1]]])

【讨论】：