向量化索引另一個數組的值數組

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]