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

vision
#1

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?

#2

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

1 Like
#3

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

#4

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?

#5

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)
    return torch.argsort(o + d)

but it does not actually meet my demands. :sweat_smile:

#6

I am going to use the following code first and see if it works: :innocent:

def rand_perm(n, K):
    o = torch.arange(0, n).float()
    d = (torch.rand_like(o) - 0.5) * K
    return torch.argsort(o + d)