# How to quickly generate a random permutation moving each element with a distance less than K?

I know the following PyTorch API can perform a global random shuffle for 1D array [0, … , n-1]:

``````torch.randperm(n)
``````

but I’m confused on how to quickly generate a random permutation, such that each element of the shuffled array satisfying:

``````K = 10  # should be positive
shuffled_array = rand_perm_func(n, K)  # a mysterious function
for i in range(n):
print(abs(shuffled_array[i] - i) < K)  # should be True for each i
``````

which means that each element is moved with a distance less than K. Does there exists a fast implementation for 1D arrays?

hm, perhaps:
r = arange(n).expand(num_arrays,-1)
scores = r + randint_like(r, high=k)
r2 = argsort(scores, dim=-1)

1 Like

It seems that this problem is very difficult… I just read this webpage: python - Controlling distance of shuffling - Stack Overflow. The following code works:

``````import torch

# randomly produces a 1-D permutation index array,
# such that each element of the shuffled array has
# a distance less than K from its original location
def rand_perm(n, K):
o = torch.arange(0, n)
if K <= 1:
return o
while True:
p = torch.randperm(n)
d = abs(p - o) < K
if bool(d.all()):
return p

if __name__ == '__main__':
for i in range(10):
print(rand_perm(10, 2))
``````

but it seems that when n is large, and K is small, the generation will take a very long time. Does there exists a more efficient implementation?

The following code is inspired by @googlebot:

``````def rand_perm(n, K):
o = torch.arange(0, n)
d = torch.randint_like(o, low=-K, high=K + 1)
but it does not actually meet my demands. I am going to use the following code first and see if it works: ``````def rand_perm(n, K):