如何以 p 的概率交换数组中的元素？

Question

Assuming I have a numpy array a = np.random.randint(0,20,10) I would like to permutate its elements with a probability p , ie if p = 0.2 each element has a 20% probability to exchange position with another element. I am aware of the numpy.random.permutate() function but this only allows to permutate all elements in an array. Can this be done efficiently in python?

Answer 1

The trick is to first choose which elements will be candidates to participate in the permutation.

import numpy as np

a = np.random.randint(0,20,10) # original array
p = 0.2

ix = np.arange(a.size)  # indexes of a
will_swap = np.random.random(a.size) <= p  # draw which elements will be candidates for swapping
after_swap = np.random.permutation(ix[will_swap]) # permute the canidadates
ix[will_swap] = after_swap # update ix with the swapped candidates
a = a[ix]
print(a) # --> [0 1 8 3 4 5 6 7 2 9]