如何快速生成随机排列移动距离小于K的每个元素？

Question

我知道以下 PyTorch API 可以对一维数组 [0, ... , n-1] 执行全局随机洗牌：

torch.randperm(n)

但我对如何快速生成随机排列感到困惑，这样打乱后的数组的每个元素都满足：

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

这意味着每个元素移动的距离小于K。一维数组和二维数组是否存在快速实现？

感谢@PM 2Ring，我写了以下代码：

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))

不过好像n大，k小的时候，生成时间会很长。有没有更高效的实现方式？

Answer 1

这是一个简单的递归生成器 Python（即不使用 PyTorch 或 Numpy），它生成满足给定约束的 range(n) 的排列。

首先，我们创建一个列表 out 来包含输出序列，将 out 中的每个插槽设置为 -1 以指示该插槽未使用。然后，对于每个值 i，我们创建一个列表 avail，其中包含允许范围内尚未被占用的索引。对于 avail 中的每个 j，我们设置 out[j] = i 并递归放置下一个 i。当 i == n 时，所有 i 都已放置，因此我们到达了递归的末尾，并且 out 包含一个有效的解决方案，该解决方案将传播回递归树。

from random import shuffle

def rand_perm_gen(n, k):
    out = [-1] * n
    def f(i):
        if i == n:
            yield out
            return

        lo, hi = max(0, i-k+1), min(n, i+k)
        avail = [j for j in range(lo, hi) if out[j] == -1]
        if not avail:
            return
        shuffle(avail)
        for j in avail:
            out[j] = i
            yield from f(i+1)
            out[j] = -1

    yield from f(0)

def test(n=10, k=3, numtests=10):
    for j, a in enumerate(rand_perm_gen(n, k), 1):
        print("\n", j, a)
        for i, u in enumerate(a):
            print(f"{i}: {u} -> {(u - i)}")
        if j == numtests:
            break

test()

典型输出


 1 [1, 0, 3, 2, 6, 5, 4, 8, 9, 7]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 8 -> 1
8: 9 -> 1
9: 7 -> -2

 2 [1, 0, 3, 2, 6, 5, 4, 9, 8, 7]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 9 -> 2
8: 8 -> 0
9: 7 -> -2

 3 [1, 0, 3, 2, 6, 5, 4, 9, 7, 8]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 9 -> 2
8: 7 -> -1
9: 8 -> -1

 4 [1, 0, 3, 2, 6, 5, 4, 8, 7, 9]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 8 -> 1
8: 7 -> -1
9: 9 -> 0

 5 [1, 0, 3, 2, 6, 5, 4, 7, 8, 9]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 7 -> 0
8: 8 -> 0
9: 9 -> 0

 6 [1, 0, 3, 2, 6, 5, 4, 7, 9, 8]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 5 -> 0
6: 4 -> -2
7: 7 -> 0
8: 9 -> 1
9: 8 -> -1

 7 [1, 0, 3, 2, 6, 7, 4, 5, 9, 8]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 7 -> 2
6: 4 -> -2
7: 5 -> -2
8: 9 -> 1
9: 8 -> -1

 8 [1, 0, 3, 2, 6, 7, 4, 5, 8, 9]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 6 -> 2
5: 7 -> 2
6: 4 -> -2
7: 5 -> -2
8: 8 -> 0
9: 9 -> 0

 9 [1, 0, 3, 2, 5, 6, 4, 7, 9, 8]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 5 -> 1
5: 6 -> 1
6: 4 -> -2
7: 7 -> 0
8: 9 -> 1
9: 8 -> -1

 10 [1, 0, 3, 2, 5, 6, 4, 7, 8, 9]
0: 1 -> 1
1: 0 -> -1
2: 3 -> 1
3: 2 -> -1
4: 5 -> 1
5: 6 -> 1
6: 4 -> -2
7: 7 -> 0
8: 8 -> 0
9: 9 -> 0

这是 SageMathCell 上的 live version 运行。

这种方法比生成所有排列并过滤它们要快，但对于大型 n 来说仍然很慢。您可以通过删除 shuffle 调用来提高速度，在这种情况下，产生的排列按字典顺序排列。

如果您只想要一个解决方案，请使用next，例如

perm = next(rand_perm_gen(10, 3))

请注意，所有解决方案共享相同的 out 列表。因此，如果您需要将这些解决方案保存在列表中，则必须复制它们，例如

perms = [seq.copy() for seq in rand_perm_gen(5, 2)]

如何快速生成随机排列移动距离小于K的每个元素？

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

python

arrays

permutation

pytorch

典型输出