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?

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?