矢量化值数组索引另一个数组

Question

我正在尝试矢量化下面的代码。

 for x in range (0, 500):
      S = rand.choice(np.unique(Y))   
      A = rand.choice(np.unique(X[Y==S])) 

Y 和 X 是索引中的值需要匹配的数组。

目前我将 S 修改为

 S = np.random.choice(np.unique(Y),size=500)

但是，我无法弄清楚如何使用值 S 数组通过 X 进行索引

大小 = 5 的示例可以是

 Y = [0,0,2,3,2,4]
 X = [1,2,1,3,4,2]
 S = [0,2,0,3,2]
 X[Y==S] => ([1,2],[1,4],[1,2],[3],[1,4])  <=  Not sure how to get this
 A = [2,4,1,3,1]

有没有简单的方法可以做到这一点？

Answer 1

有点笨拙，但如果你可以使用熊猫，isin()与列表理解相结合：

import numpy as np
import pandas as pd
Y = pd.Series([0,0,2,3,2,4])
X = pd.Series([1,2,1,3,4,2])
S = pd.Series([0,2,0,3,2])
[np.random.choice(X[Y.isin([s])].values) for s in S]

您可以通过以下方式获得X[Y==S]：[X[Y.isin([s])].values for s in S]