Reputation: 1
I am trying to vectorize the code below.
for x in range (0, 500):
S = rand.choice(np.unique(Y))
A = rand.choice(np.unique(X[Y==S]))
Y and X are arrays where the values in the indices need to match.
Currently I modified S to be
S = np.random.choice(np.unique(Y),size=500)
However, I'm not able to figure out how to index through X with an array of values S
An example for size = 5 could be
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]
Is there a simple way to do this?
Upvotes: 0
Views: 40
Reputation: 21274
A little clumsy, but if you can use Pandas, combine isin()
with a list comprehension:
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]
You can get X[Y==S]
with: [X[Y.isin([s])].values for s in S]
Upvotes: 1