python - 你如何在 numpy 数组的每一行中广播 np.random.choice ？

Question

假设我有这个 numpy 数组:

[[1, 2, 3, 4],
 [5, 6, 7, 8],
 [9, 10, 11, 12],
 [13, 14, 15, 16]]

我的目标是从每一行中选择两个随机元素并创建一个新的 numpy 数组，它可能如下所示:

[[2, 4],
 [5, 8],
 [9, 10],
 [15, 16]]

我可以使用 for 循环轻松地做到这一点。但是，有没有一种方法可以使用广播，比如 np.random.choice , 以避免必须遍历每一行？

Answer 1

方法#1
基于 this trick ，这是一种矢量化的方式 -

n = 2 # number of elements to select per row
idx = np.random.rand(*a.shape).argsort(1)[:,:n]
out = np.take_along_axis(a, idx, axis=1)

sample 运行 -

In [251]: a
Out[251]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16]])

In [252]: idx = np.random.rand(*a.shape).argsort(1)[:,:2]

In [253]: np.take_along_axis(a, idx, axis=1)
Out[253]: 
array([[ 2,  1],
       [ 6,  7],
       [ 9, 11],
       [16, 15]])

方法#2
另一个基于掩码每行选择两个 -

def select_two_per_row(a):
    m,n = a.shape
    mask = np.zeros((m,n), dtype=bool)
    R = np.arange(m)
    
    idx1 = np.random.randint(0,n,m)
    mask[R,idx1] = 1
    
    mask2 = np.zeros(m*(n-1), dtype=bool)
    idx2 = np.random.randint(0,n-1,m) + np.arange(m)*(n-1)
    mask2[idx2] = 1
    mask[~mask] = mask2
    out = a[mask].reshape(-1,2)
    return out

方法#3
另一个基于整数的索引再次选择每行两个 -

def select_two_per_row_v2(a):
    m,n = a.shape
    idx1 = np.random.randint(0,n,m)
    idx2 = np.random.randint(1,n,m)
    out = np.take_along_axis(a, np.c_[idx1, idx1 - idx2], axis=1)
    return out

时间——

In [209]: a = np.random.rand(100000,10)

# App1 with argsort
In [210]: %%timeit
     ...: idx = np.random.rand(*a.shape).argsort(1)[:,:2]
     ...: out = np.take_along_axis(a, idx, axis=1)
23.2 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# App1 with argpartition
In [221]: %%timeit
     ...: idx = np.random.rand(*a.shape).argpartition(axis=1,kth=1)[:,:2]
     ...: out = np.take_along_axis(a, idx, axis=1)
18.3 ms ± 115 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [214]: %timeit select_two_per_row(a)
9.89 ms ± 37.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [215]: %timeit select_two_per_row_v2(a)
5.78 ms ± 9.19 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)