如何选择numpy数组中行的索引?

如果您的预期输出是中每个唯一值的随机选择索引列表 y_train :

idx = [np.random.choice(np.where(y_train == i)[0], size=2, \
       replace=False) for i in np.unique(y_train)]

输出:

[array([7, 8]), array([5, 4]), array([1, 0])]

如果要将数组展平为单个数组:

idx = np.array(idx).flatten()

输出:

array([7, 8, 5, 4, 1, 0])

获得所需索引的另一种解决方案是使用 nonzero 简单的循环 range(n+1)

y_train = np.array([2,2,1,0,1,1,2,0,0])

indices = [np.random.choice((y_train==i).nonzero()[0], 2, replace=False) for i in range(n+1)]
print (indices)
# [array([7, 3]), array([5, 4]), array([0, 1])]

print (np.array(indices).ravel())
# [7 3 5 4 0 1]