Python 是否有更快的方法可以在一维数组中将一个0随机交换为一个1

发布于06月09日

我想随机交换一个一维数组中的一个0和一个1，这个数组只包含0和1，很多次都是N = 10^6.这是我做一次交换的代码.

import numpy as np

# a contains only 0 and 1
a = np.array([0,1,1,0,1,1,0,0,0,1])
# j1 random position of 0, j2 random position of 1
j1 = np.random.choice(np.where(a==0)[0])
j2 = np.random.choice(np.where(a==1)[0])
# swap
a[j1], a[j2] = a[j2], a[j1]

由于我想多次执行此过程，所以在每次迭代中，我都需要使用np.where()来定位0和1的posistion，我认为这不是很有效.

Is there any other approach which can be more effecient ?

推荐答案

在执行交换时，您可以自己维护where个表达式的结果，这样就不需要重新执行那where个调用.在程序开始时，您只需执行一次，但在循环中，您需要进行增量调整.

此外，交换可以由两个赋值代替，因为您知道0将变为1，反之亦然.

a = np.array([0,1,1,0,1,1,0,0,0,1])

# One-time preprocessing

zeroes = np.where(a==0)[0]
ones = np.where(a==1)[0]
num_zeroes = len(zeroes)
num_ones = len(ones)

# In a loop:

for _ in range(100):
    # Choose random indices in zeroes/ones
    j0 = np.random.randint(num_zeroes)
    j1 = np.random.randint(num_ones)
    # Get the chosen indices for use in `a`:
    i0 = zeroes[j0]
    i1 = ones[j1]
    # Keep the two collections in sync with the swap that now happens
    zeroes[j0] = i1
    ones[j1] = i0
    # The swap itself
    a[i0] = 1
    a[i1] = 0

    # ...

请注意，zeroes和ones的长度没有改变:我们可以重用每次交换时 Select 的两个插槽.