Index multiple dimensions of a multi-dimensional array with another - NumPy/ Python

Question

Lets say I have an tensor of the following form:

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

Now what I would like to do is to use indices as a mask on a to replace every max element with say -1. Is there a numpy way of doing this ? so far all I have figured out is using for loops.

Answer 1

Here's one way using linear indexing in 3D -

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

Sample run -

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

---

Here's another way with linear indexing again, but in 2D -

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

---

Runtime tests and verify output -

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