How to compute the inverse of permuation in numpy

By permutation, we mean a one to one and onto map from a finite set to itself. Say the set contains n elements and labeled from 0 to n-1, we can denote a permutation by a length n array a with unique elements in {0,1,...,n-1} such that a[i] is the result of i.

In numpy we can generate such permutation by shuffling [0,1,...,n-1](the identity map). Such as

rng = np.random.RandomState(0)
x = np.arange(10)
rng.shuffle(x)
print(x)

with output [2 8 4 9 1 6 7 3 0 5] which maps 0 -> 2, 1->8, etc.

It's easy to see the composition of a and b is a[b] in numpy notation. You an check

a[b][i] = a[b[i]] 

for any i in [0,1,...,n-1]

A function to compute the inverse is

def inverse_permutation(a):
    b = np.arange(a.shape[0])
    b[a] = b.copy()
    return b

It's correct because for the final b

b[a][i] = b[a[i]] = i

for any i in [0,1,...,n-1]. Note b.copy() is required or the content of b will be changed sequentially,

You can check the result by

print(x)
y = inverse_permutation(x)
print(y)
print(x[y])

with output

[2 8 4 9 1 6 7 3 0 5]
[8 4 0 7 2 9 5 6 1 3]
[0 1 2 3 4 5 6 7 8 9]