您如何在 numpy 数组的每一行中广播 np.random.choice？答案

【问题标题】：How do you broadcast np.random.choice across each row of a numpy array?您如何在 numpy 数组的每一行中广播 np.random.choice？
【发布时间】：2020-09-29 18:29:31
【问题描述】：

假设我有这个 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，以避免遍历每一行？

【问题讨论】：

使用np.apply_along_axis怎么样？

标签： python numpy

【解决方案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)

示例运行 -

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)

【讨论】：

我认为在方法 3 中，您可以使用 (idx2 - idx1) 左右保存模数。
@MadPhysicist 是的，好点。那里得到了一些改善。谢谢！
@Divakar 是方法 3 中idx1 - idx2 的分布吗？我怀疑这可能是根据与均匀随机不同的分布。
@Ehsan Well idx1 具有均匀概率选择每行 n 个元素中的任何一个。然后 idx2 在每行剩余的 n-1 个元素中具有相同的值。所以，我认为这很好。是什么让你怀疑？
@Divakar 当idx1 和idx2 都一致时，它们的减法idx1-idx2 不再一致。如果我理解正确的话，这将是两个制服的卷积，看起来像三角形分布，因此选择是不均匀的。

【解决方案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]]

【讨论】：