从每一行中随机选择（加权）其中一列的值

Question

我有大量 2 列的行。
从每一行中，我想从第一列或第二列中选择一个值（具有加权概率）。

import numpy as np

values = np.array(
    [[0.41, 0.31],
     [0.73, 0.15],
     [0.44, 0.30],
     [0.67, 0.18],
])

我想使用 0 和 1 之间的随机选择作为这样的索引，第一列的权重为 0.6，第二列的权重为 0.4：

probs_chosen = np.random.choice([0,1], size=4, replace=True, p=[0.6, 0.4])
print(probs_chosen)
array([0, 0, 1, 0])

但是如何使用该索引从第 1 行中选择第一个值，从第 2 行中选择第一个值等。
或者任何其他解决我问题的方法也可以。熊猫解决方案也可以。

这种情况下的预期结果：

[0.41, 0.73, 0.30, 0.67]

Answer 1

你可以使用 numpy advanced indexing:

row_idx = np.arange(values.shape[0])
col_idx = np.random.choice([0,1], size=4, replace=True, p=[0.6, 0.4])
out = values[row_idx, col_idx]

输出：

array([0.41, 0.73, 0.3 , 0.67])