选择关于概率向量的numpy矩阵的索引

Question

在[0,1] ^ n中给出n×n（随机）Numpy数组A和另一个Numpy数组p。对于A_i的每一行A，我想计算最小的索引j*，以便p_i <= A_i,j*。

如何在Numpy中有效实现这一点？我想这可以用numpy.random.choice以某种方式完成。

Answer 1

使用broadcasting -

的一种方法

(p.T <= A).argmax(1)

如果我们找不到任何元素p_i <= A_i,j*，我们可以使用无效的说明符，比如-1。同样，我们需要修改一个 -

mask = (p.T <= A)
out = np.where(mask.any(1), mask.argmax(1), -1)

示例运行 -

In [140]: A
Out[140]: 
array([[5, 3, 8, 0, 1],
       [5, 4, 5, 2, 6],
       [2, 5, 5, 0, 4],
       [4, 2, 6, 5, 8],
       [4, 2, 5, 2, 6]])

In [141]: p
Out[141]: array([[8, 5, 8, 5, 6]])

In [142]: mask = (p.T <= A)

In [143]: np.where(mask.any(1), mask.argmax(1), -1)
Out[143]: array([ 2,  0, -1,  2,  4])