python-3.x - 基于另一个数组的 Numpy 选择

Question

main_array如果索引对应True于另一个数组中的值的索引，我想从数组中选择一些元素。例如y应该[14,15,16]以任意顺序包含

import numpy as np

main_array = np.array([11,12,13,14,15,16])
selector = np.array([0,1,2,3,3,3])
x = np.random.choice(main_array, 3, replace=False)  # This works
y = np.random.choice(main_array, 3, replace=False, p=np.where(selector>2)) # This fails

但是，我得到ValueError: 'p' must be 1-dimensional 了将选择限制为基于另一个数组的索引的正确方法是什么？

标签： python-3.xnumpy

Answer 1

一种方法是按部分制作：

import numpy as np

main_array = np.array([11, 12, 13, 14, 15, 16])
selector = np.array([0, 1, 2, 3, 3, 3])
x = np.random.choice(main_array, 3, replace=False)
z = main_array[selector > 2]
y = np.random.choice(z, len(z), replace=False)
print(f"x={x}")
print(f"z={z}")
print(f"y={y}")

输出是

x=[16 14 13]
z=[14 15 16]
y=[16 15 14]

mask另一种方法是将不适用的概率设为零：

import numpy as np

main_array = np.array([11, 12, 13, 14, 15, 16])
selector = np.array([0, 1, 2, 3, 3, 3])
x = np.random.choice(main_array, 3, replace=False)

p = 1 * (selector > 2)
y = np.random.choice(main_array, 3, replace=False, p=p / np.sum(p))
print(y)