Efficiently define an implicit Numpy array

Question

A and B are Numpy arrays of common shape [n1,n2,n3]. The values of B are all integers in [0,n3). I want A to "invert" B in the sense that each value of A satisfies A[i,j,B[i,j,k]]=k for all i,j,k in the appropriate ranges. While it's obvious how to do this with for loops, I suspect that there is a clever one-liner using fancy indexing. Does anyone see it?

I think your question implies that for each `i` and `j`, the 1d slice `B[i,j,:]` is a permutation of `range(n3)`, otherwise the "inverse" doesn't exist. Is that correct? — Warren Weckesser, Jun 29 '17 at 02:37

Warren Weckesser · Accepted Answer · 2017-06-29T06:06:42.453

Here are two methods.

The first method is a one-liner: A = B.argsort(axis=-1)

Here's an example. B has shape (3, 5, 7) and for each fixed i and j, B[i,j,:] is a permutation of range(B.shape[2]).

In [386]: B
Out[386]: 
array([[[1, 5, 4, 6, 2, 3, 0],
        [6, 5, 3, 4, 2, 1, 0],
        [4, 5, 0, 3, 1, 2, 6],
        [0, 5, 6, 3, 2, 1, 4],
        [4, 1, 5, 2, 6, 3, 0]],

       [[2, 6, 0, 1, 5, 4, 3],
        [3, 2, 4, 0, 1, 5, 6],
        [3, 4, 6, 5, 1, 2, 0],
        [4, 6, 3, 0, 2, 5, 1],
        [0, 3, 1, 6, 4, 5, 2]],

       [[0, 3, 6, 2, 1, 5, 4],
        [3, 1, 2, 4, 6, 0, 5],
        [1, 3, 5, 6, 4, 0, 2],
        [4, 1, 6, 0, 2, 3, 5],
        [6, 4, 5, 1, 0, 3, 2]]])

In [387]: A = B.argsort(axis=-1)

In [388]: A
Out[388]: 
array([[[6, 0, 4, 5, 2, 1, 3],
        [6, 5, 4, 2, 3, 1, 0],
        [2, 4, 5, 3, 0, 1, 6],
        [0, 5, 4, 3, 6, 1, 2],
        [6, 1, 3, 5, 0, 2, 4]],

       [[2, 3, 0, 6, 5, 4, 1],
        [3, 4, 1, 0, 2, 5, 6],
        [6, 4, 5, 0, 1, 3, 2],
        [3, 6, 4, 2, 0, 5, 1],
        [0, 2, 6, 1, 4, 5, 3]],

       [[0, 4, 3, 1, 6, 5, 2],
        [5, 1, 2, 0, 3, 6, 4],
        [5, 0, 6, 1, 4, 2, 3],
        [3, 1, 4, 5, 0, 6, 2],
        [4, 3, 6, 5, 1, 2, 0]]])

Verify the desired property by sampling a few values.

In [389]: A[0, 0, B[0, 0, 0]]
Out[389]: 0

In [390]: A[0, 0, B[0, 0, 1]]
Out[390]: 1

In [391]: A[0, 0, B[0, 0, :]]
Out[391]: array([0, 1, 2, 3, 4, 5, 6])

In [392]: A[2, 3, B[2, 3, :]]
Out[392]: array([0, 1, 2, 3, 4, 5, 6])

The second method has a lower time complexity than using argsort, but it is a three-liner rather than a one-liner. I'll use the same B as above.

Create A, but with no values assigned yet.

In [393]: A = np.empty_like(B)

Create index arrays for each dimension of B.

In [394]: i, j, k = np.ogrid[[slice(n) for n in B.shape]]  # or np.ix_(*[range(n) for n in B.shape])

This is the cool part. Do the assignment exactly as you wrote it in the question.

In [395]: A[i, j, B[i, j, k]] = k

Verify that we have the same A as above.

In [396]: A
Out[396]: 
array([[[6, 0, 4, 5, 2, 1, 3],
        [6, 5, 4, 2, 3, 1, 0],
        [2, 4, 5, 3, 0, 1, 6],
        [0, 5, 4, 3, 6, 1, 2],
        [6, 1, 3, 5, 0, 2, 4]],

       [[2, 3, 0, 6, 5, 4, 1],
        [3, 4, 1, 0, 2, 5, 6],
        [6, 4, 5, 0, 1, 3, 2],
        [3, 6, 4, 2, 0, 5, 1],
        [0, 2, 6, 1, 4, 5, 3]],

       [[0, 4, 3, 1, 6, 5, 2],
        [5, 1, 2, 0, 3, 6, 4],
        [5, 0, 6, 1, 4, 2, 3],
        [3, 1, 4, 5, 0, 6, 2],
        [4, 3, 6, 5, 1, 2, 0]]])

After poking around some more on SO, I see that both these methods appear in answers to the question "How to invert a permutation array in numpy". The only thing really new here is doing the inversion along one axis of a three-dimensional array.

Efficiently define an implicit Numpy array

1 Answers1