如何在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，避免必须逐行循环？

- user4793385

---

使用np.apply_along_axis怎么样？ - deadshot

Answer 1

方法一

基于这个技巧，下面是一种向量化的方法 -

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)

示例运行 -

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

另一种基于掩码(mask)的方法是选择每行恰好两个元素 -

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

方法三

再次基于整数索引，选择每行恰好两个 -

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)

- Divakar

---

我认为在第三种方法中，你可以通过(idx2 - idx1)来节省一次模运算。 - Mad Physicist

@MadPhysicist 嗯，说得好。有一些改进的空间。谢谢！ - Divakar

@Divakar，第三种方法中idx1 - idx2的分布是否均匀？我怀疑这可能是根据不同于均匀随机分布的分布进行的。 - Ehsan

@Divakar 当idx1和idx2均匀分布时，它们的减法idx1-idx2就不再是均匀分布。这将是两个均匀分布的卷积，看起来像三角形分布，因此如果我理解正确，选择就不再是均匀的了。 - Ehsan

@Ehsan 你可能更了解分布。我的想法很简单，就是确保每行覆盖所有组合的概率相等。我看到它已经被覆盖了。这里有一些数据-https://textuploader.com/1pl44 - Divakar

Answer 2

你可以使用numpy的apply_along_axis。

import numpy as np
x = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12],  [13, 14, 15, 16]])
print(np.apply_along_axis(np.random.choice, axis=1, arr=x, size=2))

输出：

[[ 4  1]
 [ 5  6]
 [10 12]
 [14 16]]

- Ajay Verma