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?
1 Answer
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.
iandj, the 1d sliceB[i,j,:]is a permutation ofrange(n3), otherwise the "inverse" doesn't exist. Is that correct?