向量化索引另一个数组的值数组

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) 

然而，我无法通过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()：[X[Y.isin([s])].values for s in S]

：

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]