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

Question

可以说我有以下形式的张量：

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替换每个max元素。 有这样一种笨拙的方式吗？ 到目前为止，我所知道的是使用for循环。

Answer 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

Answer 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]]])