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?